Add gpu marker and test both Classic/dask-expr Dask DataFrames (#1341)

.github/workflows/test.yaml

@@ -34,7 +34,6 @@ env:
   VECLIB_MAXIMUM_THREADS: 1
   NUMEXPR_NUM_THREADS: 1
   PYDEVD_DISABLE_FILE_VALIDATION: 1
-  DASK_DATAFRAME__QUERY_PLANNING: false
 
 jobs:
   pre_commit:
@@ -174,6 +173,8 @@ jobs:
         env:
           NUMBA_DISABLE_JIT: 1
       - name: doit test_examples
+        env:
+          DASK_DATAFRAME__QUERY_PLANNING: false
         run: |
           conda activate test-environment
           doit test_examples

datashader/data_libraries/dask.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from contextlib import suppress
+
 import numpy as np
 import pandas as pd
 import dask
@@ -30,7 +32,6 @@ def _dask_compat(df):
     return getattr(df, 'optimize', lambda: df)()
 
 
-@bypixel.pipeline.register(dd.DataFrame)
 def dask_pipeline(df, schema, canvas, glyph, summary, *, antialias=False, cuda=False):
     dsk, name = glyph_dispatch(glyph, df, schema, canvas, summary, antialias=antialias, cuda=cuda)
 
@@ -50,6 +51,15 @@ def dask_pipeline(df, schema, canvas, glyph, summary, *, antialias=False, cuda=F
     return scheduler(dsk, name)
 
 
+# Classic Dask.DataFrame
+bypixel.pipeline.register(dd.core.DataFrame)(dask_pipeline)
+
+with suppress(ImportError):
+    import dask_expr
+
+    bypixel.pipeline.register(dask_expr.DataFrame)(dask_pipeline)
+
+
 def shape_bounds_st_and_axis(df, canvas, glyph):
     if not canvas.x_range or not canvas.y_range:
         x_extents, y_extents = glyph.compute_bounds_dask(df)

datashader/tests/benchmarks/test_canvas.py

@@ -1,12 +1,9 @@
 import pytest
-import os
 import numpy as np
 import pandas as pd
 
 import datashader as ds
 
-test_gpu = bool(int(os.getenv("DATASHADER_TEST_GPU", 0)))
-
 
 @pytest.fixture
 def time_series():
@@ -33,7 +30,7 @@ def test_points(benchmark, time_series):
     benchmark(cvs.points, time_series, 'x', 'y')
 
 
-@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
+@pytest.mark.gpu
 @pytest.mark.benchmark(group="canvas")
 def test_line_gpu(benchmark, time_series):
     from cudf import from_pandas
@@ -42,7 +39,7 @@ def test_line_gpu(benchmark, time_series):
     benchmark(cvs.line, time_series, 'x', 'y')
 
 
-@pytest.mark.skipif(not test_gpu, reason="DATASHADER_TEST_GPU not set")
+@pytest.mark.gpu
 @pytest.mark.benchmark(group="canvas")
 def test_points_gpu(benchmark, time_series):
     from cudf import from_pandas

datashader/tests/conftest.py

@@ -1,8 +1,8 @@
-CUSTOM_MARKS = ("benchmark",)
+CUSTOM_MARKS = {"benchmark", "gpu"}
 
 
 def pytest_addoption(parser):
-    for marker in CUSTOM_MARKS:
+    for marker in sorted(CUSTOM_MARKS):
         parser.addoption(
             f"--{marker}",
             action="store_true",
@@ -12,20 +12,21 @@ def pytest_addoption(parser):
 
 
 def pytest_configure(config):
-    for marker in CUSTOM_MARKS:
+    for marker in sorted(CUSTOM_MARKS):
         config.addinivalue_line("markers", f"{marker}: {marker} test marker")
 
 
 def pytest_collection_modifyitems(config, items):
     skipped, selected = [], []
-    markers = [m for m in CUSTOM_MARKS if config.getoption(f"--{m}")]
+    markers = {m for m in CUSTOM_MARKS if config.getoption(f"--{m}")}
     empty = not markers
     for item in items:
-        if empty and any(m in item.keywords for m in CUSTOM_MARKS):
+        item_marks = set(item.keywords) & CUSTOM_MARKS
+        if empty and item_marks:
             skipped.append(item)
         elif empty:
             selected.append(item)
-        elif not empty and any(m in item.keywords for m in markers):
+        elif not empty and item_marks == markers:
             selected.append(item)
         else:
             skipped.append(item)