add to top level import

fjetter · fjetter · commit e06101f9f599 · 2023-02-23T11:26:15.000+01:00
diff --git a/distributed/shuffle/__init__.py b/distributed/shuffle/__init__.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+from distributed.shuffle._arrow import check_minimal_arrow_version
 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__ = [
+    "check_minimal_arrow_version",
     "P2PShuffleLayer",
     "rearrange_by_column_p2p",
     "ShuffleSchedulerExtension",
diff --git a/distributed/shuffle/_arrow.py b/distributed/shuffle/_arrow.py
@@ -29,6 +29,12 @@ def check_dtype_support(meta_input: pd.DataFrame) -> None:
 
 
 def check_minimal_arrow_version() -> None:
+    """Verify that the the correct version of pyarrow is installed to support
+    the P2P extension.
+
+    Raises a RuntimeError in case pyarrow is not installed or installed version
+    is not recent enough.
+    """
     suffix = ""
     # First version to introduce Table.sort_by
     minversion = "7.0.0"
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
@@ -7,8 +7,7 @@
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import SimpleShuffleLayer
 
-from distributed.shuffle._arrow import check_dtype_support
-from distributed.shuffle._arrow import check_minimal_arrow_version
+from distributed.shuffle._arrow import check_dtype_support, check_minimal_arrow_version
 
 logger = logging.getLogger("distributed.shuffle")
 if TYPE_CHECKING:
