Skip to content

[CELEBORN-2248] Implement lazy loading for columnar shuffle classes and skew shuffle method using static holder pattern #3581

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ever4Kenny wants to merge 1 commit into
base: main
Choose a base branch
from
+53 −33
Open

[CELEBORN-2248] Implement lazy loading for columnar shuffle classes and skew shuffle method using static holder pattern #3581

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 53 additions & 33 deletions client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SparkUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,17 +228,24 @@ public static <K, C> ShuffleReader<K, C> getReader(

public static final String COLUMNAR_HASH_BASED_SHUFFLE_WRITER_CLASS =
"org.apache.spark.shuffle.celeborn.ColumnarHashBasedShuffleWriter";
static final DynConstructors.Builder COLUMNAR_HASH_BASED_SHUFFLE_WRITER_CONSTRUCTOR_BUILDER =
DynConstructors.builder()
.impl(
COLUMNAR_HASH_BASED_SHUFFLE_WRITER_CLASS,
int.class,
CelebornShuffleHandle.class,
TaskContext.class,
CelebornConf.class,
ShuffleClient.class,
ShuffleWriteMetricsReporter.class,
SendBufferPool.class);

/**
* Lazy holder for ColumnarHashBasedShuffleWriter constructor builder. The builder is initialized
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The JavaDoc comment describes a "constructor builder" but should describe a "constructor" for consistency with the recommended implementation. After storing the built DynConstructors.Ctor instead of the builder (as suggested in the comment for lines 237-247), update this comment to read "Lazy holder for ColumnarHashBasedShuffleWriter constructor. The constructor is initialized..." to accurately reflect what is being stored.

Suggested change
* Lazy holder for ColumnarHashBasedShuffleWriter constructor builder. The builder is initialized
* Lazy holder for ColumnarHashBasedShuffleWriter constructor. The constructor is initialized

Copilot uses AI. Check for mistakes.
* only when this class is first accessed, ensuring lazy loading without explicit synchronization.
*/
private static class ColumnarHashBasedShuffleWriterConstructorBuilderHolder {
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The holder class is named ColumnarHashBasedShuffleWriterConstructorBuilderHolder, but it actually holds a DynConstructors.Builder instance, not a constructor. Consider renaming to ColumnarHashBasedShuffleWriterBuilderHolder for clarity and accuracy, as "ConstructorBuilder" is redundant when the type already indicates it's a Builder for constructors.

Copilot uses AI. Check for mistakes.
private static final DynConstructors.Builder INSTANCE =
DynConstructors.builder()
.impl(
COLUMNAR_HASH_BASED_SHUFFLE_WRITER_CLASS,
int.class,
CelebornShuffleHandle.class,
TaskContext.class,
CelebornConf.class,
ShuffleClient.class,
ShuffleWriteMetricsReporter.class,
SendBufferPool.class);
Comment on lines +237 to +247
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The holder should store the built DynConstructors.Ctor instead of the DynConstructors.Builder. This is inconsistent with the established pattern in this file where all other dynamic reflection fields store the built result (see GET_READER_METHOD at line 176, UnregisterAllMapAndMergeOutput_METHOD at line 312, and even CelebornSkewShuffleMethodHolder at line 563).

Change INSTANCE to private static final DynConstructors.Ctor<?> INSTANCE and append .build() at the end of line 247. Then update line 258 to remove the .build() call and invoke directly on INSTANCE. This avoids the unnecessary method call overhead on every invocation and follows the established convention.

Copilot uses AI. Check for mistakes.
}

public static <K, V, C> HashBasedShuffleWriter<K, V, C> createColumnarHashBasedShuffleWriter(
int shuffleId,
Expand All @@ -248,26 +255,33 @@ public static <K, V, C> HashBasedShuffleWriter<K, V, C> createColumnarHashBasedS
ShuffleClient client,
ShuffleWriteMetricsReporter metrics,
SendBufferPool sendBufferPool) {
return COLUMNAR_HASH_BASED_SHUFFLE_WRITER_CONSTRUCTOR_BUILDER
return ColumnarHashBasedShuffleWriterConstructorBuilderHolder.INSTANCE
.build()
.invoke(null, shuffleId, handle, taskContext, conf, client, metrics, sendBufferPool);
}

public static final String COLUMNAR_SHUFFLE_READER_CLASS =
"org.apache.spark.shuffle.celeborn.CelebornColumnarShuffleReader";
static final DynConstructors.Builder COLUMNAR_SHUFFLE_READER_CONSTRUCTOR_BUILDER =
DynConstructors.builder()
.impl(
COLUMNAR_SHUFFLE_READER_CLASS,
CelebornShuffleHandle.class,
int.class,
int.class,
int.class,
int.class,
TaskContext.class,
CelebornConf.class,
ShuffleReadMetricsReporter.class,
ExecutorShuffleIdTracker.class);

/**
* Lazy holder for CelebornColumnarShuffleReader constructor builder. The builder is initialized
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The JavaDoc comment describes a "constructor builder" but should describe a "constructor" for consistency with the recommended implementation. After storing the built DynConstructors.Ctor instead of the builder (as suggested in the comment for lines 271-283), update this comment to read "Lazy holder for CelebornColumnarShuffleReader constructor. The constructor is initialized..." to accurately reflect what is being stored.

Suggested change
* Lazy holder for CelebornColumnarShuffleReader constructor builder. The builder is initialized
* Lazy holder for CelebornColumnarShuffleReader constructor. The constructor is initialized

Copilot uses AI. Check for mistakes.
* only when this class is first accessed, ensuring lazy loading without explicit synchronization.
*/
private static class ColumnarShuffleReaderConstructorBuilderHolder {
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The holder class is named ColumnarShuffleReaderConstructorBuilderHolder, but it actually holds a DynConstructors.Builder instance, not a constructor. Consider renaming to ColumnarShuffleReaderBuilderHolder for clarity and accuracy, as "ConstructorBuilder" is redundant when the type already indicates it's a Builder for constructors.

Copilot uses AI. Check for mistakes.
private static final DynConstructors.Builder INSTANCE =
DynConstructors.builder()
.impl(
COLUMNAR_SHUFFLE_READER_CLASS,
CelebornShuffleHandle.class,
int.class,
int.class,
int.class,
int.class,
TaskContext.class,
CelebornConf.class,
ShuffleReadMetricsReporter.class,
ExecutorShuffleIdTracker.class);
Comment on lines +271 to +283
Copy link

Copilot AI Feb 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The holder should store the built DynConstructors.Ctor instead of the DynConstructors.Builder. This is inconsistent with the established pattern in this file where all other dynamic reflection fields store the built result (see GET_READER_METHOD at line 176, UnregisterAllMapAndMergeOutput_METHOD at line 312, and even CelebornSkewShuffleMethodHolder at line 563).

Change INSTANCE to private static final DynConstructors.Ctor<?> INSTANCE and append .build() at the end of line 283. Then update line 296 to remove the .build() call and invoke directly on INSTANCE. This avoids the unnecessary method call overhead on every invocation and follows the established convention.

Copilot uses AI. Check for mistakes.
}

public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader(
CelebornShuffleHandle<K, ?, C> handle,
Expand All @@ -279,7 +293,7 @@ public static <K, C> CelebornShuffleReader<K, C> createColumnarShuffleReader(
CelebornConf conf,
ShuffleReadMetricsReporter metrics,
ExecutorShuffleIdTracker shuffleIdTracker) {
return COLUMNAR_SHUFFLE_READER_CONSTRUCTOR_BUILDER
return ColumnarShuffleReaderConstructorBuilderHolder.INSTANCE
.build()
.invoke(
null,
Expand Down Expand Up @@ -541,15 +555,21 @@ public static void addSparkListener(SparkListener listener) {
}
}

private static final DynMethods.UnboundMethod isCelebornSkewShuffle_METHOD =
DynMethods.builder("isCelebornSkewedShuffle")
.hiddenImpl("org.apache.spark.celeborn.CelebornShuffleState", Integer.TYPE)
.orNoop()
.build();
/**
* Lazy holder for isCelebornSkewedShuffle method. The method is initialized only when this class
* is first accessed, ensuring lazy loading without explicit synchronization.
*/
private static class CelebornSkewShuffleMethodHolder {
private static final DynMethods.UnboundMethod INSTANCE =
DynMethods.builder("isCelebornSkewedShuffle")
.hiddenImpl("org.apache.spark.celeborn.CelebornShuffleState", Integer.TYPE)
.orNoop()
.build();
}

public static boolean isCelebornSkewShuffleOrChildShuffle(int appShuffleId) {
if (!isCelebornSkewShuffle_METHOD.isNoop()) {
return isCelebornSkewShuffle_METHOD.asStatic().invoke(appShuffleId);
if (!CelebornSkewShuffleMethodHolder.INSTANCE.isNoop()) {
return CelebornSkewShuffleMethodHolder.INSTANCE.asStatic().invoke(appShuffleId);
} else {
return false;
}
Expand Down
Loading