Merge pull request #59 from martindurant/no_dask_import

Don't require explicit import on dask workers
intake · Jun 13, 2024 · 3f4a5ca · 3f4a5ca
2 parents eb4ff65 + a2b0be4
commit 3f4a5ca
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/awkward_pandas/dask.py b/src/awkward_pandas/dask.py
@@ -44,7 +44,7 @@ def run(self, *args, **kwargs):
                 # meta = PandasAwkwardAccessor._to_output(op(orig.ak.array, *ar, **kwargs))
                 meta = None
 
-            def inner(data):
+            def inner(data, _=DaskAwkwardAccessor):
                 import awkward_pandas.pandas  # noqa: F401
 
                 ar2 = (ar.ak.array if hasattr(ar, "ak") else ar for ar in args)

diff --git a/tests/test_dask.py b/tests/test_dask.py
@@ -36,9 +36,7 @@ def test_accessor():
 def test_distributed():
     distributed = pytest.importorskip("distributed")
 
-    with distributed.Client(
-        n_workers=1, threads_per_worker=1, preload=["awkward_pandas.dask"]
-    ):
+    with distributed.Client(n_workers=1, threads_per_worker=1):
         data = pd.arrays.ArrowExtensionArray(pa.array([[0], [0, 1]] * 2))
         s = pd.Series(data)
         df = pd.DataFrame({"s": s})