From 4243db715bfd64a1c1d405bc6f6db31507b1ef3c Mon Sep 17 00:00:00 2001 From: rjzamora Date: Wed, 24 Apr 2024 12:24:11 -0700 Subject: [PATCH] avoid using to/from_dask_dataframe --- dask_sql/input_utils/dask.py | 5 ++--- tests/integration/fixtures.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/dask_sql/input_utils/dask.py b/dask_sql/input_utils/dask.py index 4dbeeef63..006d8b6fa 100644 --- a/dask_sql/input_utils/dask.py +++ b/dask_sql/input_utils/dask.py @@ -23,11 +23,10 @@ def to_dc( ): if gpu: # pragma: no cover try: - import dask_cudf + import dask_cudf # noqa: F841 except ImportError: raise ModuleNotFoundError( "Setting `gpu=True` for table creation requires dask_cudf" ) - if not isinstance(input_item, dask_cudf.DataFrame): - input_item = dask_cudf.from_dask_dataframe(input_item, **kwargs) + return input_item.to_backend("cudf", **kwargs) return input_item diff --git a/tests/integration/fixtures.py b/tests/integration/fixtures.py index 669260215..cc6d4c142 100644 --- a/tests/integration/fixtures.py +++ b/tests/integration/fixtures.py @@ -188,7 +188,7 @@ def gpu_datetime_table(datetime_table): @pytest.fixture() def gpu_timeseries(timeseries): - return dask_cudf.from_dask_dataframe(timeseries) if dask_cudf else None + return timeseries.to_backend("cudf") if dask_cudf else None @pytest.fixture()