From b4ff54fad58f11125d40ee5d9a70349957003f11 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Wed, 12 Feb 2025 12:45:49 -0600 Subject: [PATCH] Compatibility with Dask `main` (#17992) Moves the import of get_collection_type. It's been available from `dask.dataframe` since at least `dask>=2024.12.1` xref https://github.com/dask/dask/pull/11689 and the CI failure at https://github.com/rapidsai/cudf/actions/runs/13286927884/job/37100164792?pr=17929 --- python/dask_cudf/dask_cudf/_expr/__init__.py | 167 ++++++++---------- .../dask_cudf/dask_cudf/_expr/collection.py | 2 +- 2 files changed, 76 insertions(+), 93 deletions(-) diff --git a/python/dask_cudf/dask_cudf/_expr/__init__.py b/python/dask_cudf/dask_cudf/_expr/__init__.py index 329ad861081..1f757476ce5 100644 --- a/python/dask_cudf/dask_cudf/_expr/__init__.py +++ b/python/dask_cudf/dask_cudf/_expr/__init__.py @@ -1,96 +1,79 @@ # Copyright (c) 2024-2025, NVIDIA CORPORATION. -from packaging.version import Version - import dask +import dask.dataframe.dask_expr._shuffle as _shuffle_module +from dask.dataframe import get_collection_type +from dask.dataframe.dask_expr import ( + DataFrame as DXDataFrame, + FrameBase, + Index as DXIndex, + Series as DXSeries, + from_dict, + new_collection, +) +from dask.dataframe.dask_expr._cumulative import ( + CumulativeBlockwise, +) +from dask.dataframe.dask_expr._expr import ( + Elemwise, + Expr, + RenameAxis, + VarColumns, +) +from dask.dataframe.dask_expr._groupby import ( + DecomposableGroupbyAggregation, + GroupBy as DXGroupBy, + GroupbyAggregation, + SeriesGroupBy as DXSeriesGroupBy, + SingleAggregation, +) +from dask.dataframe.dask_expr._reductions import ( + Reduction, + Var, +) +from dask.dataframe.dask_expr._util import ( + _convert_to_list, + _raise_if_object_series, + is_scalar, +) +from dask.dataframe.dask_expr.io.io import ( + FusedIO, + FusedParquetIO, +) +from dask.dataframe.dask_expr.io.parquet import ( + FragmentWrapper, + ReadParquetFSSpec, + ReadParquetPyarrowFS, +) -if Version(dask.__version__) > Version("2024.12.1"): - import dask.dataframe.dask_expr._shuffle as _shuffle_module - from dask.dataframe.dask_expr import ( - DataFrame as DXDataFrame, - FrameBase, - Index as DXIndex, - Series as DXSeries, - from_dict, - get_collection_type, - new_collection, - ) - from dask.dataframe.dask_expr._cumulative import ( - CumulativeBlockwise, - ) - from dask.dataframe.dask_expr._expr import ( - Elemwise, - Expr, - RenameAxis, - VarColumns, - ) - from dask.dataframe.dask_expr._groupby import ( - DecomposableGroupbyAggregation, - GroupBy as DXGroupBy, - GroupbyAggregation, - SeriesGroupBy as DXSeriesGroupBy, - SingleAggregation, - ) - from dask.dataframe.dask_expr._reductions import ( - Reduction, - Var, - ) - from dask.dataframe.dask_expr._util import ( - _convert_to_list, - _raise_if_object_series, - is_scalar, - ) - from dask.dataframe.dask_expr.io.io import ( - FusedIO, - FusedParquetIO, - ) - from dask.dataframe.dask_expr.io.parquet import ( - FragmentWrapper, - ReadParquetFSSpec, - ReadParquetPyarrowFS, - ) -else: - import dask_expr._shuffle as _shuffle_module # noqa: F401 - from dask_expr import ( # noqa: F401 - DataFrame as DXDataFrame, - FrameBase, - Index as DXIndex, - Series as DXSeries, - from_dict, - get_collection_type, - new_collection, - ) - from dask_expr._cumulative import CumulativeBlockwise # noqa: F401 - from dask_expr._expr import ( # noqa: F401 - Elemwise, - Expr, - RenameAxis, - VarColumns, - ) - from dask_expr._groupby import ( # noqa: F401 - DecomposableGroupbyAggregation, - GroupBy as DXGroupBy, - GroupbyAggregation, - SeriesGroupBy as DXSeriesGroupBy, - SingleAggregation, - ) - from dask_expr._reductions import Reduction, Var # noqa: F401 - from dask_expr._util import ( # noqa: F401 - _convert_to_list, - _raise_if_object_series, - is_scalar, - ) - from dask_expr.io.io import FusedIO, FusedParquetIO # noqa: F401 - from dask_expr.io.parquet import ( # noqa: F401 - FragmentWrapper, - ReadParquetFSSpec, - ReadParquetPyarrowFS, - ) - - from dask.dataframe import _dask_expr_enabled - - if not _dask_expr_enabled(): - raise ValueError( - "The legacy DataFrame API is not supported for RAPIDS >24.12. " - "The 'dataframe.query-planning' config must be True or None." - ) +__all__ = [ + "CumulativeBlockwise", + "DXDataFrame", + "DXGroupBy", + "DXIndex", + "DXSeries", + "DXSeriesGroupBy", + "DecomposableGroupbyAggregation", + "Elemwise", + "Expr", + "FragmentWrapper", + "FrameBase", + "FusedIO", + "FusedParquetIO", + "GroupbyAggregation", + "ReadParquetFSSpec", + "ReadParquetPyarrowFS", + "Reduction", + "RenameAxis", + "SingleAggregation", + "Var", + "VarColumns", + "_convert_to_list", + "_raise_if_object_series", + "_shuffle_module", + "dask", + "from_dict", + "get_collection_type", + "is_scalar", + "new_collection", +] diff --git a/python/dask_cudf/dask_cudf/_expr/collection.py b/python/dask_cudf/dask_cudf/_expr/collection.py index c3e44567abf..9e42ea2d650 100644 --- a/python/dask_cudf/dask_cudf/_expr/collection.py +++ b/python/dask_cudf/dask_cudf/_expr/collection.py @@ -4,6 +4,7 @@ from functools import cached_property from dask import config +from dask.dataframe import get_collection_type from dask.dataframe.core import is_dataframe_like from dask.dataframe.dispatch import get_parallel_type from dask.typing import no_default @@ -16,7 +17,6 @@ DXSeries, FrameBase, _raise_if_object_series, - get_collection_type, new_collection, )