This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch branch-0.7
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/branch-0.7 by this push:
new 99f135af51 [CELEBORN-2480] Push columnar shuffle batches by estimated
size in bytes
99f135af51 is described below
commit 99f135af51f65dcd3e6f5cfbc5256474619beff1
Author: ForVic <[email protected]>
AuthorDate: Tue Oct 20 20:37:12 2026 +0800
[CELEBORN-2480] Push columnar shuffle batches by estimated size in bytes
### What changes were proposed in this pull request?
Add an optional byte limit to the columnar shuffle writer. When
`celeborn.columnarShuffle.batch.sizeInBytes.enabled` is true,
`ColumnarHashBasedShuffleWriter` pushes a partition's batch once it reaches
either `celeborn.columnarShuffle.batch.size` rows or
`celeborn.columnarShuffle.batch.sizeInBytes.threshold` estimated bytes,
whichever comes first. The estimate is the sum of `UnsafeRow.getSizeInBytes()`
over the rows in the batch, so the column builders are unchanged.
| Config | Default |
|---|---|
| `celeborn.columnarShuffle.batch.sizeInBytes.enabled` | `false` |
| `celeborn.columnarShuffle.batch.sizeInBytes.threshold` | `64k`, same as
`celeborn.client.push.buffer.max.size` |
The same change is made to the spark-3, spark-3.5 and spark-4 columnar
writers.
### Why are the changes needed?
The columnar writer keeps one in-memory batch per reduce partition and
pushes it only after 10,000 rows. Row count is a poor proxy for memory: 10,000
rows can be a few hundred KB or hundreds of MB, depending on string widths.
Spark's memory manager does not track this memory, so it cannot spill.
If a map task has fewer than 10,000 rows per partition, nothing is pushed
until the task ends. That is common: 1M rows over 200 partitions is 5,000 rows
each. The task's whole output then stays on the heap. The row-based hash writer
instead caps each partition's buffer at `celeborn.client.push.buffer.max.size`,
which defaults to 64 KiB.
Example: repartition 4M rows with a 2 KiB string into 200 partitions, using
4 map tasks:
```scala
spark.range(0, 4000000, 1, 4)
.selectExpr("id", "cast(id % 1000 as int) as category",
"repeat(md5(cast(id as string)), 64) as payload")
.repartition(200, $"id")
.write.format("noop").mode("overwrite").save()
```
Smallest driver heap that completes the job (tried 1g, 2g, 4g, 8g, 16g):
| Writer | Min heap |
|---|---|
| Row-based hash writer | 1g |
| Columnar, row limit only (current behavior) | OOM even at 16g |
| Columnar + byte limit, `64k` (default) | 1g |
| Columnar + byte limit, `256k` | 2g |
| Columnar + byte limit, `1m` | 4g |
| Columnar + byte limit, `4m` | 8g |
As with the row-based writer, buffered row data is roughly `partitions x
threshold`. This is not a hard memory cap: as before this change, the column
builders pre-size their buffers for `celeborn.columnarShuffle.batch.size` rows
(up to 80 KB per column per partition). Sizing those buffers from the threshold
could be a follow-up.
Same job at a 24g heap, where every configuration fits (mean of 2 runs):
| Columnar | Wall | Task time | Shuffle bytes |
|---|---|---|---|
| Row limit only | 12.3s | 40.7s | 213.9 MB |
| `64k` | 9.8s | 31.4s | 230.7 MB (+7.9%) |
| `256k` | 8.0s | 25.3s | 221.8 MB (+3.7%) |
| `1m` | 8.3s | 26.3s | 218.3 MB (+2.1%) |
Smaller batches compress slightly worse, but they avoid most of the GC
work. Narrow rows are barely affected: 40M rows of 64 bytes, 50,000 rows per
partition per task, 4g heap:
| Columnar | Wall | Shuffle bytes |
|---|---|---|
| Row limit only | 11.4s | 1655.9 MB |
| `64k` | 11.5s | 1594.7 MB |
| `256k` | 10.8s | 1663.0 MB |
| `1m` | 11.0s | 1655.9 MB (byte limit never reached) |
The default of `64k` matches the row-based writer's per-partition buffer
size. A larger threshold trades memory for slightly smaller shuffle output.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
Adds the two configs above. Behavior is unchanged unless
`celeborn.columnarShuffle.batch.sizeInBytes.enabled` is set.
### How was this patch tested?
- Added
`ColumnarHashBasedShuffleWriterSuiteJ#testPushBatchWhenSizeInBytesThresholdReached`
to the spark-3, spark-3.5 and spark-4 columnar modules. It checks the rows in
each pushed batch, with the byte limit (8, 8, 8, 6) and without it (30).
- Ran `build/mvn -P<profile> -pl client-spark/<module> -am test
-Dtest=ColumnarHashBasedShuffleWriterSuiteJ -DwildcardSuites=none` for
`spark-3.3`, `spark-3.5` and `spark-4.0`. All 14 tests pass in each. Maven CI
doesn't run the columnar modules, so these were run locally.
- Regenerated `docs/configuration/columnar-shuffle.md` with
`ConfigurationSuite`.
- Ran the jobs above on Spark 3.5.8 in `local[4]` with one local worker,
JDK 17, on a 16 vCPU arm64 host. Spark had
`assets/spark-patch/Celeborn_Columnar_Shuffle_spark3.patch` applied. On 3.5 the
patch needs `DataTypeUtils.fromAttributes` in place of
`StructType.fromAttributes` to compile. Timings are from a shared machine, so
treat small differences as noise.
Closes #3873 from ForVic/columnar-shuffle-batch-bytes.
Authored-by: ForVic <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
(cherry picked from commit 659587e321836a061893b302996aba416c9190e7)
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../celeborn/ColumnarHashBasedShuffleWriter.java | 12 +++-
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 71 ++++++++++++++++++++++
.../celeborn/ColumnarHashBasedShuffleWriter.java | 12 +++-
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 71 ++++++++++++++++++++++
.../celeborn/ColumnarHashBasedShuffleWriter.java | 12 +++-
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 71 ++++++++++++++++++++++
.../org/apache/celeborn/common/CelebornConf.scala | 29 +++++++++
docs/configuration/columnar-shuffle.md | 2 +
docs/developers/java-columnar-shuffle.md | 6 ++
9 files changed, 283 insertions(+), 3 deletions(-)
diff --git
a/client-spark/spark-3-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
b/client-spark/spark-3-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
index b09b1306c9..225eeaa91d 100644
---
a/client-spark/spark-3-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
+++
b/client-spark/spark-3-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
@@ -53,6 +53,8 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
private final Serializer depSerializer;
private final boolean isColumnarShuffle;
private final int columnarShuffleBatchSize;
+ private final long columnarShuffleBatchSizeInBytes;
+ private final long[] celebornBatchBytes;
private final boolean columnarShuffleCodeGenEnabled;
private final boolean columnarShuffleDictionaryEnabled;
private final double columnarShuffleDictionaryMaxFactor;
@@ -68,6 +70,10 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
throws IOException {
super(shuffleId, handle, taskContext, conf, client, metrics,
sendBufferPool);
columnarShuffleBatchSize = conf.columnarShuffleBatchSize();
+ columnarShuffleBatchSizeInBytes =
+ conf.columnarShuffleBatchSizeInBytesEnabled()
+ ? conf.columnarShuffleBatchSizeInBytesThreshold()
+ : Long.MAX_VALUE;
columnarShuffleCodeGenEnabled = conf.columnarShuffleCodeGenEnabled();
columnarShuffleDictionaryEnabled = conf.columnarShuffleDictionaryEnabled();
columnarShuffleDictionaryMaxFactor =
conf.columnarShuffleDictionaryMaxFactor();
@@ -78,6 +84,7 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
this.depSerializer = handle.dependency().serializer();
this.celebornBatchBuilders =
new
CelebornBatchBuilder[handle.dependency().partitioner().numPartitions()];
+ this.celebornBatchBytes = new long[celebornBatchBuilders.length];
this.isColumnarShuffle = schema != null &&
CelebornBatchBuilder.supportsColumnarType(schema);
}
@@ -119,13 +126,16 @@ public class ColumnarHashBasedShuffleWriter<K, V, C>
extends HashBasedShuffleWri
}
celebornBatchBuilders[partitionId].writeRow(row);
- if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize) {
+ celebornBatchBytes[partitionId] += row.getSizeInBytes();
+ if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize
+ || celebornBatchBytes[partitionId] >=
columnarShuffleBatchSizeInBytes) {
byte[] arr = celebornBatchBuilders[partitionId].buildColumnBytes();
pushGiantRecord(partitionId, arr, arr.length);
if (dataSize != null) {
dataSize.add(arr.length);
}
celebornBatchBuilders[partitionId].newBuilders();
+ celebornBatchBytes[partitionId] = 0;
}
tmpRecordsWritten++;
}
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 92aee552fb..4cc24575c4 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
@@ -17,25 +17,38 @@
package org.apache.spark.shuffle.celeborn;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import java.io.File;
+import java.nio.ByteBuffer;
import java.util.UUID;
+import scala.Product2;
+import scala.Tuple2;
+import scala.collection.Iterator;
+import scala.collection.mutable.ListBuffer;
+
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.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
import org.apache.spark.shuffle.ShuffleWriter;
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
import org.apache.spark.sql.execution.UnsafeRowSerializer;
import org.apache.spark.sql.execution.columnar.CelebornBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchSerializer;
import org.apache.spark.sql.types.IntegerType$;
import org.apache.spark.sql.types.StringType$;
import org.apache.spark.sql.types.StructType;
+import org.apache.spark.unsafe.types.UTF8String;
import org.junit.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
@@ -43,6 +56,8 @@ import org.mockito.Mockito;
import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.reflect.DynConstructors;
public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterSuiteBase {
@@ -83,6 +98,62 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
assertTrue(((ColumnarHashBasedShuffleWriter<?, ?, ?>)
writer).isColumnarShuffle());
}
+ @Test
+ public void testPushBatchWhenSizeInBytesThresholdReached() throws Exception {
+ // A (key, 100-byte value) UnsafeRow takes 8 + 2 * 8 + 104 = 128 bytes, so
a 1k threshold
+ // pushes a batch every 8 rows.
+ 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");
+ assertArrayEquals(new int[] {8, 8, 8, 6}, writeAndGetBatchRowCounts(conf,
30));
+
+ // Without it, only the row limit applies and all rows stay in one batch.
+ assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+ }
+
+ private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ final String partitionIdPassthroughClazz =
+ SparkVersionUtil.isGreaterThan(3, 3)
+ ? "org.apache.spark.PartitionIdPassthrough"
+ : "org.apache.spark.sql.execution.PartitionIdPassthrough";
+ final Partitioner partitioner =
+ DynConstructors.builder()
+ .impl(partitionIdPassthroughClazz, int.class)
+ .<Partitioner>build()
+ .newInstance(numPartitions);
+ Mockito.doReturn(partitioner).when(dependency).partitioner();
+ Mockito.doReturn(new UnsafeRowSerializer(2,
null)).when(dependency).serializer();
+
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
+ final ShuffleWriter<Integer, String> writer =
+ createShuffleWriter(
+ new CelebornShuffleHandle<>(
+ appId, host, port, userIdentifier, shuffleId, false, numMaps,
dependency),
+ taskContext,
+ conf,
+ client,
+ metrics.shuffleWriteMetrics());
+
+ final UnsafeProjection projection = UnsafeProjection.create(schema);
+ final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ 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()));
+ }
+ 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();
+ }
+
@Override
protected SerializerInstance newSerializerInstance(Serializer serializer) {
if (serializer instanceof UnsafeRowSerializer
diff --git
a/client-spark/spark-3.5-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
b/client-spark/spark-3.5-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
index b09b1306c9..225eeaa91d 100644
---
a/client-spark/spark-3.5-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
+++
b/client-spark/spark-3.5-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
@@ -53,6 +53,8 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
private final Serializer depSerializer;
private final boolean isColumnarShuffle;
private final int columnarShuffleBatchSize;
+ private final long columnarShuffleBatchSizeInBytes;
+ private final long[] celebornBatchBytes;
private final boolean columnarShuffleCodeGenEnabled;
private final boolean columnarShuffleDictionaryEnabled;
private final double columnarShuffleDictionaryMaxFactor;
@@ -68,6 +70,10 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
throws IOException {
super(shuffleId, handle, taskContext, conf, client, metrics,
sendBufferPool);
columnarShuffleBatchSize = conf.columnarShuffleBatchSize();
+ columnarShuffleBatchSizeInBytes =
+ conf.columnarShuffleBatchSizeInBytesEnabled()
+ ? conf.columnarShuffleBatchSizeInBytesThreshold()
+ : Long.MAX_VALUE;
columnarShuffleCodeGenEnabled = conf.columnarShuffleCodeGenEnabled();
columnarShuffleDictionaryEnabled = conf.columnarShuffleDictionaryEnabled();
columnarShuffleDictionaryMaxFactor =
conf.columnarShuffleDictionaryMaxFactor();
@@ -78,6 +84,7 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
this.depSerializer = handle.dependency().serializer();
this.celebornBatchBuilders =
new
CelebornBatchBuilder[handle.dependency().partitioner().numPartitions()];
+ this.celebornBatchBytes = new long[celebornBatchBuilders.length];
this.isColumnarShuffle = schema != null &&
CelebornBatchBuilder.supportsColumnarType(schema);
}
@@ -119,13 +126,16 @@ public class ColumnarHashBasedShuffleWriter<K, V, C>
extends HashBasedShuffleWri
}
celebornBatchBuilders[partitionId].writeRow(row);
- if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize) {
+ celebornBatchBytes[partitionId] += row.getSizeInBytes();
+ if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize
+ || celebornBatchBytes[partitionId] >=
columnarShuffleBatchSizeInBytes) {
byte[] arr = celebornBatchBuilders[partitionId].buildColumnBytes();
pushGiantRecord(partitionId, arr, arr.length);
if (dataSize != null) {
dataSize.add(arr.length);
}
celebornBatchBuilders[partitionId].newBuilders();
+ celebornBatchBytes[partitionId] = 0;
}
tmpRecordsWritten++;
}
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 92aee552fb..4cc24575c4 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
@@ -17,25 +17,38 @@
package org.apache.spark.shuffle.celeborn;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import java.io.File;
+import java.nio.ByteBuffer;
import java.util.UUID;
+import scala.Product2;
+import scala.Tuple2;
+import scala.collection.Iterator;
+import scala.collection.mutable.ListBuffer;
+
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.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
import org.apache.spark.shuffle.ShuffleWriter;
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
import org.apache.spark.sql.execution.UnsafeRowSerializer;
import org.apache.spark.sql.execution.columnar.CelebornBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchSerializer;
import org.apache.spark.sql.types.IntegerType$;
import org.apache.spark.sql.types.StringType$;
import org.apache.spark.sql.types.StructType;
+import org.apache.spark.unsafe.types.UTF8String;
import org.junit.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
@@ -43,6 +56,8 @@ import org.mockito.Mockito;
import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.reflect.DynConstructors;
public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterSuiteBase {
@@ -83,6 +98,62 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
assertTrue(((ColumnarHashBasedShuffleWriter<?, ?, ?>)
writer).isColumnarShuffle());
}
+ @Test
+ public void testPushBatchWhenSizeInBytesThresholdReached() throws Exception {
+ // A (key, 100-byte value) UnsafeRow takes 8 + 2 * 8 + 104 = 128 bytes, so
a 1k threshold
+ // pushes a batch every 8 rows.
+ 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");
+ assertArrayEquals(new int[] {8, 8, 8, 6}, writeAndGetBatchRowCounts(conf,
30));
+
+ // Without it, only the row limit applies and all rows stay in one batch.
+ assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+ }
+
+ private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ final String partitionIdPassthroughClazz =
+ SparkVersionUtil.isGreaterThan(3, 3)
+ ? "org.apache.spark.PartitionIdPassthrough"
+ : "org.apache.spark.sql.execution.PartitionIdPassthrough";
+ final Partitioner partitioner =
+ DynConstructors.builder()
+ .impl(partitionIdPassthroughClazz, int.class)
+ .<Partitioner>build()
+ .newInstance(numPartitions);
+ Mockito.doReturn(partitioner).when(dependency).partitioner();
+ Mockito.doReturn(new UnsafeRowSerializer(2,
null)).when(dependency).serializer();
+
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
+ final ShuffleWriter<Integer, String> writer =
+ createShuffleWriter(
+ new CelebornShuffleHandle<>(
+ appId, host, port, userIdentifier, shuffleId, false, numMaps,
dependency),
+ taskContext,
+ conf,
+ client,
+ metrics.shuffleWriteMetrics());
+
+ final UnsafeProjection projection = UnsafeProjection.create(schema);
+ final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ 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()));
+ }
+ 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();
+ }
+
@Override
protected SerializerInstance newSerializerInstance(Serializer serializer) {
if (serializer instanceof UnsafeRowSerializer
diff --git
a/client-spark/spark-4-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
b/client-spark/spark-4-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
index b09b1306c9..225eeaa91d 100644
---
a/client-spark/spark-4-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
+++
b/client-spark/spark-4-columnar-shuffle/src/main/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriter.java
@@ -53,6 +53,8 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
private final Serializer depSerializer;
private final boolean isColumnarShuffle;
private final int columnarShuffleBatchSize;
+ private final long columnarShuffleBatchSizeInBytes;
+ private final long[] celebornBatchBytes;
private final boolean columnarShuffleCodeGenEnabled;
private final boolean columnarShuffleDictionaryEnabled;
private final double columnarShuffleDictionaryMaxFactor;
@@ -68,6 +70,10 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
throws IOException {
super(shuffleId, handle, taskContext, conf, client, metrics,
sendBufferPool);
columnarShuffleBatchSize = conf.columnarShuffleBatchSize();
+ columnarShuffleBatchSizeInBytes =
+ conf.columnarShuffleBatchSizeInBytesEnabled()
+ ? conf.columnarShuffleBatchSizeInBytesThreshold()
+ : Long.MAX_VALUE;
columnarShuffleCodeGenEnabled = conf.columnarShuffleCodeGenEnabled();
columnarShuffleDictionaryEnabled = conf.columnarShuffleDictionaryEnabled();
columnarShuffleDictionaryMaxFactor =
conf.columnarShuffleDictionaryMaxFactor();
@@ -78,6 +84,7 @@ public class ColumnarHashBasedShuffleWriter<K, V, C> extends
HashBasedShuffleWri
this.depSerializer = handle.dependency().serializer();
this.celebornBatchBuilders =
new
CelebornBatchBuilder[handle.dependency().partitioner().numPartitions()];
+ this.celebornBatchBytes = new long[celebornBatchBuilders.length];
this.isColumnarShuffle = schema != null &&
CelebornBatchBuilder.supportsColumnarType(schema);
}
@@ -119,13 +126,16 @@ public class ColumnarHashBasedShuffleWriter<K, V, C>
extends HashBasedShuffleWri
}
celebornBatchBuilders[partitionId].writeRow(row);
- if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize) {
+ celebornBatchBytes[partitionId] += row.getSizeInBytes();
+ if (celebornBatchBuilders[partitionId].getRowCnt() >=
columnarShuffleBatchSize
+ || celebornBatchBytes[partitionId] >=
columnarShuffleBatchSizeInBytes) {
byte[] arr = celebornBatchBuilders[partitionId].buildColumnBytes();
pushGiantRecord(partitionId, arr, arr.length);
if (dataSize != null) {
dataSize.add(arr.length);
}
celebornBatchBuilders[partitionId].newBuilders();
+ celebornBatchBytes[partitionId] = 0;
}
tmpRecordsWritten++;
}
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 92aee552fb..4cc24575c4 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
@@ -17,25 +17,38 @@
package org.apache.spark.shuffle.celeborn;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import java.io.File;
+import java.nio.ByteBuffer;
import java.util.UUID;
+import scala.Product2;
+import scala.Tuple2;
+import scala.collection.Iterator;
+import scala.collection.mutable.ListBuffer;
+
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.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
import org.apache.spark.shuffle.ShuffleWriteMetricsReporter;
import org.apache.spark.shuffle.ShuffleWriter;
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow;
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection;
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
import org.apache.spark.sql.execution.UnsafeRowSerializer;
import org.apache.spark.sql.execution.columnar.CelebornBatchBuilder;
import org.apache.spark.sql.execution.columnar.CelebornColumnarBatchSerializer;
import org.apache.spark.sql.types.IntegerType$;
import org.apache.spark.sql.types.StringType$;
import org.apache.spark.sql.types.StructType;
+import org.apache.spark.unsafe.types.UTF8String;
import org.junit.Test;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
@@ -43,6 +56,8 @@ import org.mockito.Mockito;
import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.reflect.DynConstructors;
public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterSuiteBase {
@@ -83,6 +98,62 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
assertTrue(((ColumnarHashBasedShuffleWriter<?, ?, ?>)
writer).isColumnarShuffle());
}
+ @Test
+ public void testPushBatchWhenSizeInBytesThresholdReached() throws Exception {
+ // A (key, 100-byte value) UnsafeRow takes 8 + 2 * 8 + 104 = 128 bytes, so
a 1k threshold
+ // pushes a batch every 8 rows.
+ 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");
+ assertArrayEquals(new int[] {8, 8, 8, 6}, writeAndGetBatchRowCounts(conf,
30));
+
+ // Without it, only the row limit applies and all rows stay in one batch.
+ assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+ }
+
+ private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ final String partitionIdPassthroughClazz =
+ SparkVersionUtil.isGreaterThan(3, 3)
+ ? "org.apache.spark.PartitionIdPassthrough"
+ : "org.apache.spark.sql.execution.PartitionIdPassthrough";
+ final Partitioner partitioner =
+ DynConstructors.builder()
+ .impl(partitionIdPassthroughClazz, int.class)
+ .<Partitioner>build()
+ .newInstance(numPartitions);
+ Mockito.doReturn(partitioner).when(dependency).partitioner();
+ Mockito.doReturn(new UnsafeRowSerializer(2,
null)).when(dependency).serializer();
+
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
+ final ShuffleWriter<Integer, String> writer =
+ createShuffleWriter(
+ new CelebornShuffleHandle<>(
+ appId, host, port, userIdentifier, shuffleId, false, numMaps,
dependency),
+ taskContext,
+ conf,
+ client,
+ metrics.shuffleWriteMetrics());
+
+ final UnsafeProjection projection = UnsafeProjection.create(schema);
+ final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ 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()));
+ }
+ 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();
+ }
+
@Override
protected SerializerInstance newSerializerInstance(Serializer serializer) {
if (serializer instanceof UnsafeRowSerializer
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 1fadc77f23..fb557fee18 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -1493,6 +1493,10 @@ class CelebornConf(loadDefaults: Boolean) extends
Cloneable with Logging with Se
// //////////////////////////////////////////////////////
def columnarShuffleEnabled: Boolean = get(COLUMNAR_SHUFFLE_ENABLED)
def columnarShuffleBatchSize: Int = get(COLUMNAR_SHUFFLE_BATCH_SIZE)
+ def columnarShuffleBatchSizeInBytesEnabled: Boolean =
+ get(COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_ENABLED)
+ def columnarShuffleBatchSizeInBytesThreshold: Long =
+ get(COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_THRESHOLD)
def columnarShuffleOffHeapEnabled: Boolean =
get(COLUMNAR_SHUFFLE_OFF_HEAP_ENABLED)
def columnarShuffleDictionaryEnabled: Boolean =
get(COLUMNAR_SHUFFLE_DICTIONARY_ENCODING_ENABLED)
def columnarShuffleDictionaryMaxFactor: Double =
@@ -6266,6 +6270,31 @@ object CelebornConf extends Logging {
.checkValue(v => v > 0, "Value must be positive")
.createWithDefault(10000)
+ val COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_ENABLED: ConfigEntry[Boolean] =
+ buildConf("celeborn.columnarShuffle.batch.sizeInBytes.enabled")
+ .categories("columnar-shuffle")
+ .version("1.0.0")
+ .doc("Whether to also push a columnar batch once its estimated size in
bytes reaches " +
+ "`celeborn.columnarShuffle.batch.sizeInBytes.threshold`, in addition
to the row " +
+ s"limit `${COLUMNAR_SHUFFLE_BATCH_SIZE.key}`. This limits how much row
data each " +
+ "partition buffers when rows are wide, e.g. with long strings. Column
buffers are " +
+ s"still pre-sized for `${COLUMNAR_SHUFFLE_BATCH_SIZE.key}` rows.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_THRESHOLD: ConfigEntry[Long] =
+ buildConf("celeborn.columnarShuffle.batch.sizeInBytes.threshold")
+ .categories("columnar-shuffle")
+ .version("1.0.0")
+ .doc("Estimated size in bytes at which a columnar batch is pushed, when
" +
+ "`celeborn.columnarShuffle.batch.sizeInBytes.enabled` is true. The
size is estimated " +
+ "from the uncompressed size of the rows added to the batch. The
default matches " +
+ s"`${CLIENT_PUSH_BUFFER_MAX_SIZE.key}`, the per-partition buffer of
the row-based " +
+ "hash writer.")
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(v => v > 0, "Value must be positive")
+ .createWithDefaultString("64k")
+
val COLUMNAR_SHUFFLE_OFF_HEAP_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.columnarShuffle.offHeap.enabled")
.withAlternative("celeborn.columnar.offHeap.enabled")
diff --git a/docs/configuration/columnar-shuffle.md
b/docs/configuration/columnar-shuffle.md
index 3152dbfbea..af6fe0deab 100644
--- a/docs/configuration/columnar-shuffle.md
+++ b/docs/configuration/columnar-shuffle.md
@@ -20,6 +20,8 @@ license: |
| Key | Default | isDynamic | Description | Since | Deprecated |
| --- | ------- | --------- | ----------- | ----- | ---------- |
| celeborn.columnarShuffle.batch.size | 10000 | false | Vector batch size for
columnar shuffle. | 0.3.0 | celeborn.columnar.shuffle.batch.size |
+| celeborn.columnarShuffle.batch.sizeInBytes.enabled | false | false | Whether
to also push a columnar batch once its estimated size in bytes reaches
`celeborn.columnarShuffle.batch.sizeInBytes.threshold`, in addition to the row
limit `celeborn.columnarShuffle.batch.size`. This limits how much row data each
partition buffers when rows are wide, e.g. with long strings. Column buffers
are still pre-sized for `celeborn.columnarShuffle.batch.size` rows. | 1.0.0 |
|
+| celeborn.columnarShuffle.batch.sizeInBytes.threshold | 64k | false |
Estimated size in bytes at which a columnar batch is pushed, when
`celeborn.columnarShuffle.batch.sizeInBytes.enabled` is true. The size is
estimated from the uncompressed size of the rows added to the batch. The
default matches `celeborn.client.push.buffer.max.size`, the per-partition
buffer of the row-based hash writer. | 1.0.0 | |
| celeborn.columnarShuffle.codegen.enabled | false | false | Whether to use
codegen for columnar-based shuffle. | 0.3.0 |
celeborn.columnar.shuffle.codegen.enabled |
| celeborn.columnarShuffle.enabled | false | false | Whether to enable
columnar-based shuffle. | 0.2.0 | celeborn.columnar.shuffle.enabled |
| celeborn.columnarShuffle.encoding.dictionary.enabled | false | false |
Whether to use dictionary encoding for columnar-based shuffle data. | 0.3.0 |
celeborn.columnar.shuffle.encoding.dictionary.enabled |
diff --git a/docs/developers/java-columnar-shuffle.md
b/docs/developers/java-columnar-shuffle.md
index c01df0bbc0..40f1444a2c 100644
--- a/docs/developers/java-columnar-shuffle.md
+++ b/docs/developers/java-columnar-shuffle.md
@@ -54,3 +54,9 @@ If you require further performance optimization, consider
enabling code generati
```
spark.celeborn.columnarShuffle.codegen.enabled true
```
+
+Each partition buffers rows in memory until
`spark.celeborn.columnarShuffle.batch.size` rows are reached. If rows can be
wide, e.g. with long strings, also cap each batch by its estimated size in
bytes to limit how much data each partition buffers:
+
+```
+spark.celeborn.columnarShuffle.batch.sizeInBytes.enabled true
+```
