diff --git a/distributed/shuffle/__init__.py b/distributed/shuffle/__init__.py index e77ddacdf10..c4858b4b2df 100644 --- a/distributed/shuffle/__init__.py +++ b/distributed/shuffle/__init__.py @@ -1,10 +1,13 @@ from __future__ import annotations +from distributed.shuffle._merge import HashJoinP2PLayer, hash_join_p2p from distributed.shuffle._scheduler_extension import ShuffleSchedulerExtension from distributed.shuffle._shuffle import P2PShuffleLayer, rearrange_by_column_p2p from distributed.shuffle._worker_extension import ShuffleWorkerExtension __all__ = [ + "hash_join_p2p", + "HashJoinP2PLayer", "P2PShuffleLayer", "rearrange_by_column_p2p", "ShuffleSchedulerExtension",