fugue 0.9.6__tar.gz → 0.9.7.dev0__tar.gz

{fugue-0.9.6 → fugue-0.9.7.dev0}/PKG-INFO

@@ -1,27 +1,26 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.6
+Version: 0.9.7.dev0
 Summary: An abstraction layer for distributed computation
-Home-page: http://github.com/fugue-project/fugue
-Author: The Fugue Development Team
-Author-email: hello@fugue.ai
+Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0
-Keywords: distributed spark dask ray sql dsl domain specific language
+Project-URL: Homepage, http://github.com/fugue-project/fugue
+Project-URL: Repository, http://github.com/fugue-project/fugue
+Keywords: distributed,spark,dask,ray,sql,dsl,domain specific language
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: triad>=1.0.0
+Requires-Dist: triad>=1.0.1
 Requires-Dist: adagio>=0.2.6
 Requires-Dist: pandas<3
 Provides-Extra: sql
@@ -39,7 +38,7 @@ Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "dask"
 Requires-Dist: pyarrow>=7.0.0; extra == "dask"
 Requires-Dist: pandas>=2.0.2; extra == "dask"
 Provides-Extra: ray
-Requires-Dist: ray[data]>=2.30.0; extra == "ray"
+Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "ray"
 Requires-Dist: duckdb>=0.5.0; extra == "ray"
 Requires-Dist: pyarrow>=7.0.0; extra == "ray"
 Requires-Dist: pandas; extra == "ray"
@@ -70,7 +69,7 @@ Requires-Dist: jinja2; extra == "all"
 Requires-Dist: pyspark>=3.1.1; extra == "all"
 Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "all"
 Requires-Dist: dask-sql; extra == "all"
-Requires-Dist: ray[data]>=2.30.0; extra == "all"
+Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "all"
 Requires-Dist: notebook; extra == "all"
 Requires-Dist: jupyterlab; extra == "all"
 Requires-Dist: ipython>=7.10.0; extra == "all"
@@ -79,26 +78,14 @@ Requires-Dist: pyarrow>=6.0.1; extra == "all"
 Requires-Dist: pandas>=2.0.2; extra == "all"
 Requires-Dist: ibis-framework[duckdb,pandas]; extra == "all"
 Requires-Dist: polars; extra == "all"
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: keywords
-Dynamic: license
 Dynamic: license-file
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
 
-# Fugue
+# <img src="./images/logo.svg" width="200">
 
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/master/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/README.md

@@ -3,7 +3,7 @@
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/master/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/__init__.py

@@ -1,4 +1,8 @@
 # flake8: noqa
+from importlib.metadata import version
+
+__version__ = version("fugue")
+
 from triad.collections import Schema
 
 from fugue.api import out_transform, transform
@@ -83,7 +87,6 @@ from fugue.sql.workflow import FugueSQLWorkflow
 from fugue.workflow._workflow_context import FugueWorkflowContext
 from fugue.workflow.module import module
 from fugue.workflow.workflow import FugueWorkflow, WorkflowDataFrame, WorkflowDataFrames
-from fugue_version import __version__
 
 from .dev import *
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/dataframe/dataframe.py

@@ -59,9 +59,7 @@ class DataFrame(Dataset):
             assert isinstance(self._schema, Schema)
             return self._schema  # type: ignore
         with self._lazy_schema_lock:
-            self._schema = _input_schema(
-                self._schema()
-            ).assert_not_empty()  # type: ignore
+            self._schema = _input_schema(self._schema()).assert_not_empty()  # type: ignore
             self._schema.set_readonly()
             self._schema_discovered = True
             return self._schema

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/dataframe/function_wrapper.py

@@ -1,4 +1,5 @@
 import inspect
+from collections.abc import Callable as AbcCallable
 from typing import (
     Any,
     Callable,
@@ -7,6 +8,9 @@ from typing import (
     Iterator,
     List,
     Optional,
+    Union,
+    get_args,
+    get_origin,
     no_type_check,
 )
 
@@ -39,13 +43,66 @@ from .pandas_dataframe import PandasDataFrame
 
 
 def _compare_iter(tp: Any) -> Any:
-    return lambda x: compare_annotations(
-        x, Iterable[tp]  # type:ignore
-    ) or compare_annotations(
-        x, Iterator[tp]  # type:ignore
+    return lambda x: (
+        compare_annotations(
+            x,
+            Iterable[tp],  # type:ignore
+        )
+        or compare_annotations(
+            x,
+            Iterator[tp],  # type:ignore
+        )
     )
 
 
+def _is_optional(annotation) -> bool:
+    origin = get_origin(annotation)
+
+    # Check if it's a Union type
+    if origin is Union:
+        args = get_args(annotation)
+        # Optional[T] is Union[T, None]
+        return type(None) in args
+
+
+def _is_required_callable(annotation) -> bool:
+    """Check if annotation is a required (non-optional) Callable type."""
+    if _is_optional(annotation):
+        return False
+
+    # Check direct equality
+    if annotation == Callable or annotation == callable:  # pylint: disable=comparison-with-callable
+        return True
+
+    # Check if it's a generic Callable like Callable[[int], str]
+    origin = get_origin(annotation)
+    return origin is AbcCallable or origin is type(Callable)
+
+
+def _is_optional_callable(annotation) -> bool:
+    """Check if annotation is an optional Callable type (Optional[Callable] or Callable | None)."""
+    if not _is_optional(annotation):
+        return False
+
+    # Get the non-None types from the Union
+    args = get_args(annotation)
+    non_none_types = [arg for arg in args if arg is not type(None)]
+
+    # Should have exactly one non-None type, and it should be Callable
+    if len(non_none_types) != 1:  # pragma: no cover
+        return False
+
+    inner_type = non_none_types[0]
+
+    # Check if the inner type is Callable
+    if inner_type == Callable or inner_type == callable:  # pylint: disable=comparison-with-callable
+        return True
+
+    # Check if it's a generic Callable like Callable[[int], str]
+    origin = get_origin(inner_type)
+    return origin is AbcCallable or origin is type(Callable)
+
+
 @function_wrapper(FUGUE_ENTRYPOINT)
 class DataFrameFunctionWrapper(FunctionWrapper):
     @property
@@ -154,12 +211,7 @@ fugue_annotated_param = DataFrameFunctionWrapper.annotated_param
 @fugue_annotated_param(
     "Callable",
     "F",
-    lambda annotation: (
-        annotation == Callable
-        or annotation == callable  # pylint: disable=comparison-with-callable
-        or str(annotation).startswith("typing.Callable")
-        or str(annotation).startswith("collections.abc.Callable")
-    ),
+    _is_required_callable,
 )
 class _CallableParam(AnnotatedParam):
     pass
@@ -168,15 +220,7 @@ class _CallableParam(AnnotatedParam):
 @fugue_annotated_param(
     "Callable",
     "f",
-    lambda annotation: (
-        annotation == Optional[Callable]
-        or annotation == Optional[callable]
-        or str(annotation).startswith("typing.Union[typing.Callable")  # 3.8-
-        or str(annotation).startswith("typing.Optional[typing.Callable")  # 3.9+
-        or str(annotation).startswith(
-            "typing.Optional[collections.abc.Callable]"
-        )  # 3.9+
-    ),
+    _is_optional_callable,
 )
 class _OptionalCallableParam(AnnotatedParam):
     pass

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/dataframe/utils.py

@@ -61,9 +61,9 @@ def _df_eq(
     else:
         df2 = as_fugue_df(data).as_local_bounded()
     try:
-        assert (
-            df1.count() == df2.count()
-        ), f"count mismatch {df1.count()}, {df2.count()}"
+        assert df1.count() == df2.count(), (
+            f"count mismatch {df1.count()}, {df2.count()}"
+        )
         assert not check_schema or df.schema.is_like(
             df2.schema, equal_groups=equal_type_groups
         ), f"schema mismatch {df.schema.pa_schema}, {df2.schema.pa_schema}"

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/extensions/_builtins/outputters.py

@@ -67,7 +67,7 @@ class Save(Outputter):
             mode=mode,
             partition_spec=partition_spec,
             force_single=force_single,
-            **kwargs
+            **kwargs,
         )
 
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/extensions/_builtins/processors.py

@@ -314,7 +314,7 @@ class SaveAndUse(Processor):
             mode=mode,
             partition_spec=partition_spec,
             force_single=force_single,
-            **kwargs
+            **kwargs,
         )
         return self.execution_engine.load_df(path=path, format_hint=format_hint)
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/extensions/processor/convert.py

@@ -136,6 +136,7 @@ def processor(
     Please read
     :doc:`Processor Tutorial <tutorial:tutorials/extensions/processor>`
     """
+
     # TODO: validation of schema if without * should be done at compile time
     def deco(func: Callable) -> "_FuncAsProcessor":
         return _FuncAsProcessor.from_func(

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/extensions/transformer/convert.py

@@ -546,7 +546,10 @@ class _FuncAsOutputCoTransformer(_FuncAsCoTransformer):
             p = dict(dfs)
             p.update(self.params)
             self._wrapper.run(
-                [] + cb, p, ignore_unknown=False, output=False  # type: ignore
+                [] + cb,
+                p,
+                ignore_unknown=False,
+                output=False,  # type: ignore
             )
         return ArrayDataFrame([], OUTPUT_TRANSFORMER_DUMMY_SCHEMA)
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/sql/_utils.py

@@ -35,7 +35,6 @@ def fill_sql_template(sql: str, params: Dict[str, Any]):
         template = Template(new_sql)
 
     except jinja2.exceptions.TemplateSyntaxError:
-
         template = Template(sql)
 
     return template.render(**params)

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/sql/_visitors.py

@@ -820,7 +820,6 @@ class _Extensions(_VisitorBase):
 
     def visitSetOperation(self, ctx: fp.SetOperationContext) -> Iterable[Any]:
         def get_sub(_ctx: Tree) -> List[Any]:
-
             sub = list(
                 self.visitFugueTerm(_ctx)
                 if isinstance(_ctx, fp.FugueTermContext)

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/workflow/_tasks.py

@@ -204,7 +204,9 @@ class FugueTask(TaskSpec, ABC):
 
             # add caller traceback
             ctb = modify_traceback(
-                sys.exc_info()[2].tb_next, None, self._traceback  # type: ignore
+                sys.exc_info()[2].tb_next,
+                None,
+                self._traceback,  # type: ignore
             )
             if ctb is None:  # pragma: no cover
                 raise

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue/workflow/workflow.py

@@ -1468,7 +1468,8 @@ class WorkflowDataFrames(DataFrames):
         super().__setitem__(key, value, *args, **kwds)
 
     def __getitem__(  # pylint: disable=W0235
-        self, key: Union[str, int]  # type: ignore
+        self,
+        key: Union[str, int],  # type: ignore
     ) -> WorkflowDataFrame:
         return super().__getitem__(key)  # type: ignore
 
@@ -1791,8 +1792,7 @@ class FugueWorkflow:
                 )
             )
         raise FugueWorkflowCompileError(
-            f"Input data of type {type(data)} can't "
-            "be converted to WorkflowDataFrame"
+            f"Input data of type {type(data)} can't be converted to WorkflowDataFrame"
         )
 
     def df(

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue.egg-info/PKG-INFO

@@ -1,27 +1,26 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.6
+Version: 0.9.7.dev0
 Summary: An abstraction layer for distributed computation
-Home-page: http://github.com/fugue-project/fugue
-Author: The Fugue Development Team
-Author-email: hello@fugue.ai
+Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0
-Keywords: distributed spark dask ray sql dsl domain specific language
+Project-URL: Homepage, http://github.com/fugue-project/fugue
+Project-URL: Repository, http://github.com/fugue-project/fugue
+Keywords: distributed,spark,dask,ray,sql,dsl,domain specific language
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: triad>=1.0.0
+Requires-Dist: triad>=1.0.1
 Requires-Dist: adagio>=0.2.6
 Requires-Dist: pandas<3
 Provides-Extra: sql
@@ -39,7 +38,7 @@ Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "dask"
 Requires-Dist: pyarrow>=7.0.0; extra == "dask"
 Requires-Dist: pandas>=2.0.2; extra == "dask"
 Provides-Extra: ray
-Requires-Dist: ray[data]>=2.30.0; extra == "ray"
+Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "ray"
 Requires-Dist: duckdb>=0.5.0; extra == "ray"
 Requires-Dist: pyarrow>=7.0.0; extra == "ray"
 Requires-Dist: pandas; extra == "ray"
@@ -70,7 +69,7 @@ Requires-Dist: jinja2; extra == "all"
 Requires-Dist: pyspark>=3.1.1; extra == "all"
 Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "all"
 Requires-Dist: dask-sql; extra == "all"
-Requires-Dist: ray[data]>=2.30.0; extra == "all"
+Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "all"
 Requires-Dist: notebook; extra == "all"
 Requires-Dist: jupyterlab; extra == "all"
 Requires-Dist: ipython>=7.10.0; extra == "all"
@@ -79,26 +78,14 @@ Requires-Dist: pyarrow>=6.0.1; extra == "all"
 Requires-Dist: pandas>=2.0.2; extra == "all"
 Requires-Dist: ibis-framework[duckdb,pandas]; extra == "all"
 Requires-Dist: polars; extra == "all"
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: keywords
-Dynamic: license
 Dynamic: license-file
-Dynamic: provides-extra
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
 
-# Fugue
+# <img src="./images/logo.svg" width="200">
 
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/master/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue.egg-info/SOURCES.txt

@@ -1,7 +1,6 @@
 LICENSE
 README.md
-setup.cfg
-setup.py
+pyproject.toml
 fugue/__init__.py
 fugue/api.py
 fugue/constants.py
@@ -157,5 +156,4 @@ fugue_test/bag_suite.py
 fugue_test/builtin_suite.py
 fugue_test/dataframe_suite.py
 fugue_test/execution_suite.py
-fugue_test/fixtures.py
-fugue_version/__init__.py
+fugue_test/fixtures.py

fugue-0.9.7.dev0/fugue.egg-info/entry_points.txt

@@ -0,0 +1,11 @@
+[fugue.plugins]
+dask = fugue_dask.registry
+duckdb = fugue_duckdb.registry
+ibis = fugue_ibis
+polars = fugue_polars.registry
+ray = fugue_ray.registry
+spark = fugue_spark.registry
+
+[pytest11]
+fugue_test = fugue_test
+fugue_test_fixtures = fugue_test.fixtures

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue.egg-info/requires.txt

@@ -1,4 +1,4 @@
-triad>=1.0.0
+triad>=1.0.1
 adagio>=0.2.6
 pandas<3
 
@@ -10,7 +10,6 @@ jinja2
 pyspark>=3.1.1
 dask[dataframe,distributed]>=2024.4.0
 dask-sql
-ray[data]>=2.30.0
 notebook
 jupyterlab
 ipython>=7.10.0
@@ -20,6 +19,9 @@ pandas>=2.0.2
 ibis-framework[duckdb,pandas]
 polars
 
+[all:python_version < "3.14"]
+ray[data]>=2.30.0
+
 [cpp_sql_parser]
 fugue-sql-antlr[cpp]>=0.2.0
 
@@ -52,11 +54,13 @@ ipython>=7.10.0
 polars
 
 [ray]
-ray[data]>=2.30.0
 duckdb>=0.5.0
 pyarrow>=7.0.0
 pandas
 
+[ray:python_version < "3.14"]
+ray[data]>=2.30.0
+
 [spark]
 pyspark>=3.1.1
 zstandard>=0.25.0

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue.egg-info/top_level.txt

@@ -9,4 +9,3 @@ fugue_ray
 fugue_spark
 fugue_sql
 fugue_test
-fugue_version

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue_dask/__init__.py

@@ -1,5 +1,7 @@
 # flake8: noqa
-from fugue_version import __version__
+from importlib.metadata import version
+
+__version__ = version("fugue")
 
 from fugue_dask.dataframe import DaskDataFrame
 from fugue_dask.execution_engine import DaskExecutionEngine

{fugue-0.9.6 → fugue-0.9.7.dev0}/fugue_dask/_utils.py

@@ -55,8 +55,7 @@ def hash_repartition(df: dd.DataFrame, num: int, cols: List[Any]) -> dd.DataFram
     if num == 1:
         return df.repartition(npartitions=1)
     df = df.reset_index(drop=True).clear_divisions()
-    idf, ct = _add_hash_index(df, num, cols)
-    return _postprocess(idf, ct, num)
+    return _add_hash_index(df, num, cols)
 
 
 def even_repartition(df: dd.DataFrame, num: int, cols: List[Any]) -> dd.DataFrame:
@@ -81,13 +80,9 @@ def even_repartition(df: dd.DataFrame, num: int, cols: List[Any]) -> dd.DataFram
         return df
     df = df.reset_index(drop=True).clear_divisions()
     if len(cols) == 0:
-        idf, ct = _add_continuous_index(df)
+        return _add_continuous_index(df, num=num)
     else:
-        idf, ct = _add_group_index(df, cols, shuffle=False)
-        # when cols are set and num is not set, we use the number of groups
-        if num <= 0:
-            num = ct
-    return _postprocess(idf, ct, num)
+        return _add_group_index(df, cols, shuffle=False, num=num)
 
 
 def rand_repartition(
@@ -114,25 +109,30 @@ def rand_repartition(
         return df.repartition(npartitions=1)
     df = df.reset_index(drop=True).clear_divisions()
     if len(cols) == 0:
-        idf, ct = _add_random_index(df, num=num, seed=seed)
+        return _add_random_index(df, num=num, seed=seed)
     else:
-        idf, ct = _add_group_index(df, cols, shuffle=True, seed=seed)
-        # when cols are set and num is not set, we use the number of groups
-    return _postprocess(idf, ct, num)
+        return _add_group_index(df, cols, shuffle=True, num=num, seed=seed)
 
 
-def _postprocess(idf: dd.DataFrame, ct: int, num: int) -> dd.DataFrame:
-    parts = min(ct, num)
+def _safe_set_index(df: dd.DataFrame, key_ct: int, num_partitions: int) -> dd.DataFrame:
+    if num_partitions <= 0:
+        num_partitions = key_ct
+    parts = min(key_ct, num_partitions)
     if parts <= 1:
-        return idf.repartition(npartitions=1)
-    divisions = list(np.arange(ct, step=math.ceil(ct / parts)))
-    divisions.append(ct - 1)
-    return idf.repartition(divisions=divisions, force=True)
+        return df.set_index(
+            _FUGUE_DASK_TEMP_IDX_COLUMN, drop=True, sort=True, npartitions=1
+        )
+    divisions = np.arange(key_ct, step=int(math.ceil(key_ct / parts))).tolist()
+    # divisions.append(ct - 1)
+    divisions.append(key_ct)
+    return df.set_index(
+        _FUGUE_DASK_TEMP_IDX_COLUMN, drop=True, sort=True, divisions=divisions
+    )
 
 
 def _add_group_index(
-    df: dd.DataFrame, cols: List[str], shuffle: bool, seed: Any = None
-) -> Tuple[dd.DataFrame, int]:
+    df: dd.DataFrame, cols: List[str], shuffle: bool, num: int, seed: Any = None
+) -> dd.DataFrame:
     keys = df[cols].drop_duplicates().compute()
     if shuffle:
         keys = keys.sample(frac=1, random_state=seed)
@@ -140,12 +140,10 @@ def _add_group_index(
         **{_FUGUE_DASK_TEMP_IDX_COLUMN: pd.Series(range(len(keys)), dtype=int)}
     )
     df = df.merge(dd.from_pandas(keys, npartitions=1), on=cols, broadcast=True)
-    return df.set_index(_FUGUE_DASK_TEMP_IDX_COLUMN, drop=True), len(keys)
+    return _safe_set_index(df, len(keys), num)
 
 
-def _add_hash_index(
-    df: dd.DataFrame, num: int, cols: List[str]
-) -> Tuple[dd.DataFrame, int]:
+def _add_hash_index(df: dd.DataFrame, num: int, cols: List[str]) -> dd.DataFrame:
     if len(cols) == 0:
         cols = list(df.columns)
 
@@ -165,13 +163,13 @@ def _add_hash_index(
     orig_schema = list(df.dtypes.to_dict().items())
     idf = df.map_partitions(
         _add_hash, meta=orig_schema + [(_FUGUE_DASK_TEMP_IDX_COLUMN, int)]
-    ).set_index(_FUGUE_DASK_TEMP_IDX_COLUMN, drop=True)
-    return idf, num
+    )
+    return _safe_set_index(idf, num, num)
 
 
 def _add_random_index(
     df: dd.DataFrame, num: int, seed: Any = None
-) -> Tuple[dd.DataFrame, int]:  # pragma: no cover
+) -> dd.DataFrame:  # pragma: no cover
     def _add_rand(df: pd.DataFrame) -> pd.DataFrame:
         if len(df) == 0:
             return df.assign(**{_FUGUE_DASK_TEMP_IDX_COLUMN: pd.Series(dtype=int)})
@@ -184,11 +182,11 @@ def _add_random_index(
     orig_schema = list(df.dtypes.to_dict().items())
     idf = df.map_partitions(
         _add_rand, meta=orig_schema + [(_FUGUE_DASK_TEMP_IDX_COLUMN, int)]
-    ).set_index(_FUGUE_DASK_TEMP_IDX_COLUMN, drop=True)
-    return idf, num
+    )
+    return _safe_set_index(idf, num, num)
 
 
-def _add_continuous_index(df: dd.DataFrame) -> Tuple[dd.DataFrame, int]:
+def _add_continuous_index(df: dd.DataFrame, num: int) -> dd.DataFrame:
     def _get_info(
         df: pd.DataFrame, partition_info: Any
     ) -> pd.DataFrame:  # pragma: no cover
@@ -216,8 +214,7 @@ def _add_continuous_index(df: dd.DataFrame) -> Tuple[dd.DataFrame, int]:
     idf = df.map_partitions(
         _add_index, meta=orig_schema + [(_FUGUE_DASK_TEMP_IDX_COLUMN, int)]
     )
-    idf = idf.set_index(_FUGUE_DASK_TEMP_IDX_COLUMN, drop=True)
-    return idf, counts[-1]
+    return _safe_set_index(idf, counts[-1], num)
 
 
 class DaskUtils(PandasLikeUtils[dd.DataFrame, dd.Series]):
@@ -255,7 +252,7 @@ class DaskUtils(PandasLikeUtils[dd.DataFrame, dd.Series]):
         schema: pa.Schema,
         use_extension_types: bool = True,
         use_arrow_dtype: bool = False,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> DataFrame:
         output_dtypes = to_pandas_dtype(
             schema,
@@ -268,7 +265,7 @@ class DaskUtils(PandasLikeUtils[dd.DataFrame, dd.Series]):
             use_extension_types=use_extension_types,
             use_arrow_dtype=use_arrow_dtype,
             meta=output_dtypes,
-            **kwargs
+            **kwargs,
         )