Repository: dask-contrib/dask-expr
Pull request #878 (2024-02-14 15:54:35)

Fix divisions for head
closes https://github.com/coiled/benchmarks/issues/1393


cc @jrbourbeau 


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index c20254a..ce39a1c 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -4857,6 +4857,7 @@ def from_map(
     NOTE: The underlying ``Expr`` object produced by this API
     will support column projection (via ``simplify``) if
     the ``func`` argument has "columns" in its signature.
+
     """
     from dask.dataframe.io.utils import DataFrameIOFunction
 
@@ -4887,8 +4888,16 @@ def from_map(
 
     # Check if `func` supports column projection
     allow_projection = False
-    if "columns" in inspect.signature(func).parameters:
+    columns_arg_required = False
+    if param := inspect.signature(func).parameters.get("columns", None):
         allow_projection = True
+        columns_arg_required = param.default is param.empty
+        if meta is no_default and columns_arg_required:
+            raise TypeError(
+                "Argument `func` of `from_map` has a required `columns` "
+                " parameter and not `meta` provided."
+                "Either provide `meta` yourself or make `columns` an optional argument."
+            )
     elif isinstance(func, DataFrameIOFunction):
         warnings.warn(
             "dask_expr does not support the DataFrameIOFunction "
@@ -4910,6 +4919,7 @@ def from_map(
                 columns,
                 args,
                 kwargs,
+                columns_arg_required,
                 meta,
                 enforce_metadata,
                 divisions,
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 655d5a1..6d800ea 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -2190,7 +2190,7 @@ class Head(Expr):
             return self.frame.divisions[0], self.frame.divisions[-1]
         return (
             self.frame.divisions[0],
-            self.frame.divisions[self.operand("npartitions") + 1],
+            self.frame.divisions[self.operand("npartitions")],
         )
 
     def _task(self, index: int):
diff --git a/dask_expr/io/io.py b/dask_expr/io/io.py
index 25ef1ab..1b6b34f 100644
--- a/dask_expr/io/io.py
+++ b/dask_expr/io/io.py
@@ -228,6 +228,7 @@ class FromMapProjectable(FromMap):
         "columns",
         "args",
         "kwargs",
+        "columns_arg_required",
         "user_meta",
         "enforce_metadata",
         "user_divisions",
@@ -265,9 +266,9 @@ class FromMapProjectable(FromMap):
     @functools.cached_property
     def kwargs(self):
         options = self.operand("kwargs")
-        if self.columns_operand:
+        if self.columns_arg_required or self.columns_operand:
             options = options.copy()
-            options["columns"] = self.columns_operand
+            options["columns"] = self.columns
         return options
 
     @functools.cached_property


Test patch:
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index 3bea37f..cca5d63 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -7,7 +7,7 @@ import numpy as np
 import pyarrow.parquet as pq
 import pytest
 from dask.array.utils import assert_eq as array_assert_eq
-from dask.dataframe.utils import assert_eq
+from dask.dataframe.utils import assert_eq, make_meta
 
 from dask_expr import (
     DataFrame,
@@ -463,6 +463,29 @@ def test_from_map(tmpdir, meta, label, allow_projection, enforce_metadata):
     assert_eq(result, pdf["a"], check_index=False)
 
 
+def func(path, columns):
+    raise NotImplementedError("This shouldn't ever be called")
+
+
+def func2(path, columns):
+    return pd.DataFrame({"a": range(10), "b": range(10)})
+
+
+def test_from_map_columns_required():
+    with pytest.raises(TypeError, match=r"columns.*optional"):
+        from_map(func, ["foo"])
+    meta = make_meta(func2("foo", None))
+    ddf = from_map(func2, ["foo"], meta=meta)
+    ddf_expected = from_pandas(func2("foo", None), npartitions=1)
+
+    assert_eq(ddf, ddf_expected, check_divisions=False)
+
+    actual = from_map(func2, ["foo"], meta=meta)[["a"]].optimize()
+    expected = from_map(func2, ["foo"], meta=meta, columns=["a"]).optimize()
+
+    assert actual._name == expected._name
+
+
 def test_from_array():
     arr = np.random.randint(1, 100, (100,))
     assert_eq(from_array(arr, chunksize=5), pd.Series(arr))
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index bbdee63..3404f21 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -2336,6 +2336,11 @@ def test_shift_with_freq_datetime(pdf, data_freq, divs1):
     assert res.known_divisions == divs1
 
 
+def test_head_npartitions_divisions(pdf):
+    df = from_pandas(pdf, npartitions=1)
+    assert_eq(df.head(3, compute=False), pdf.head(3))
+
+
 @pytest.mark.parametrize("data_freq,divs", [("D", True), ("h", True)])
 def test_shift_with_freq_period_index(pdf, data_freq, divs):
     pdf.index = pd.period_range(start="2020-01-01", periods=len(pdf), freq=data_freq)