use forward- and backward-compatible logic for importing dask_expr (#477)

Signed-off-by: rjzamora <rzamora217@gmail.com>
Co-authored-by: Ayush Dattagupta <ayushdg95@gmail.com>

Author: rjzamora

nemo_curator/_compat.py

@@ -64,7 +64,9 @@ def query_planning_enabled():
     global _DASK_QUERY_PLANNING_ENABLED
 
     if _DASK_QUERY_PLANNING_ENABLED is None:
-        if _dask_version > parse_version("2024.6.0"):
+        if _dask_version > parse_version("2024.12.1"):
+            _DASK_QUERY_PLANNING_ENABLED = True
+        elif _dask_version > parse_version("2024.6.0"):
             import dask.dataframe as dd
 
             _DASK_QUERY_PLANNING_ENABLED = dd.DASK_EXPR_ENABLED

nemo_curator/utils/fuzzy_dedup_utils/shuffle_utils.py

@@ -65,12 +65,15 @@ def rearange_by_column_direct(
         )
 
     elif query_planning_enabled():
-        from dask_expr._collection import new_collection
-        from dask_expr._shuffle import RearrangeByColumn
+        try:
+            from dask.dataframe import dask_expr
+        except ImportError:
+            # TODO: Remove when pinned to dask>2024.12.1
+            import dask_expr
 
         # Use the internal dask-expr API
-        return new_collection(
-            RearrangeByColumn(
+        return dask_expr.new_collection(
+            dask_expr._shuffle.RearrangeByColumn(
                 frame=df.expr,
                 partitioning_index=col,
                 npartitions_out=npartitions,