Repository: dask/dask

Issue:
dask.bag.map(): "TypeError: apply() got an unexpected keyword argument"
Passing keyword arguments to dask.bag.map() doesn't seem to work for me:

```
def f(x, y):
    return {x: y}
testbag = db.from_sequence(['x']).map(f, y='y')
testbag.to_dataframe()

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-7387e49cb28a> in <module>()
      2     return {x: y}
      3 testbag = db.from_sequence(['x']).map(f, y='y')
----> 4 testbag.to_dataframe()

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/bag/core.py in to_dataframe(self, columns)
   1118             meta = columns
   1119         else:
-> 1120             head = self.take(1)[0]
   1121             meta = pd.DataFrame([head], columns=columns)
   1122         columns = list(meta.columns)

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/bag/core.py in take(self, k, npartitions, compute)
   1002 
   1003         if compute:
-> 1004             return tuple(b.compute())
   1005         else:
   1006             return b

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/base.py in compute(self, **kwargs)
     84             Extra keywords to forward to the scheduler ``get`` function.
     85         """
---> 86         return compute(self, **kwargs)[0]
     87 
     88     @classmethod

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/base.py in compute(*args, **kwargs)
    177         dsk = merge(var.dask for var in variables)
    178     keys = [var._keys() for var in variables]
--> 179     results = get(dsk, keys, **kwargs)
    180 
    181     results_iter = iter(results)

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/multiprocessing.py in get(dsk, keys, num_workers, func_loads, func_dumps, optimize_graph, **kwargs)
     81         # Run
     82         result = get_async(apply_async, len(pool._pool), dsk3, keys,
---> 83                            queue=queue, get_id=_process_get_id, **kwargs)
     84     finally:
     85         if cleanup:

/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/async.py in get_async(apply_async, num_workers, dsk, result, cache, queue, get_id, raise_on_exception, rerun_exceptions_locally, callbacks, **kwargs)
    479                     _execute_task(task, data)  # Re-execute locally
    480                 else:
--> 481                     raise(remote_exception(res, tb))
    482             state['cache'][key] = res
    483             finish_task(dsk, key, state, results, keyorder.get)

TypeError: apply() got an unexpected keyword argument 'y'

Traceback
---------
  File "/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/async.py", line 263, in execute_task
    result = _execute_task(task, data)
  File "/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/async.py", line 245, in _execute_task
    return func(*args2)
  File "/home_local/caichinger/miniconda3/lib/python3.5/site-packages/dask/bag/core.py", line 1667, in safe_take
    r = list(take(n, b))
```

Happens with dask 0.11.0 and conda 4.1.11, Python 3.5.2 on Ubuntu 16.04.



Discussion:
This has to do with some magical behavior in `dask.bag.Bag.map`, where functions with multiple arguments are implicitly applied using `*args`. In pseudocode:

``` python
def map(f, bag):
    if takes_multiple_arguments(f):
        return [f(*item) for item in bag]
    else:
        return [f(item) for item in bag]
```

If you make `y` a keyword argument instead of an argument, then your code works:

``` python
In [1]: import dask.bag as db

In [2]: def f(x, y=None):
   ...:     return {x: y}
   ...:

In [3]: db.from_sequence(['x']).map(f, y='y').to_dataframe().compute()
Out[3]:
   x
0  y
```

@mrocklin: perhaps it's time to reconsider my argument against this magic behavior (https://github.com/dask/dask/issues/472#issuecomment-125249022). This seems to have enough edge cases that it's hard to understand what's going on. We also don't implicitly do this for other methods either (`filter`, etc...). This has caused a few user issues:
- https://github.com/dask/dask/issues/472
- https://github.com/dask/dask/issues/1062

Alternatively, we could just informatively error when passing in kwargs to functions that take multiple arguments, since it's hard to tell whether we should splat or if the remaining arguments are filled by keywords.

I'll admit that I remain fond of the magic, but am happy to see it go if it's causing more issues than it has value.  We'll have to be careful about raising informative errors when it goes away.

The source of my problems with Bag.map():
- The actual behavior undocumented and actually contradicts the documentation "Keyword arguments are passed through to func"
- The error message (`apply() got an unexpected keyword argument 'y'`) gave no hint what I was doing wrong

The inconsistency where **kwargs are forwarded to is confusing. Also, what is the difference between the magic map behavior and map_partitions()?

I am also getting a similar error when working through the example Notebooks:

df.B[(df.A > 0) & (df.A < 3)].sum().compute()

---

TypeError                    Traceback (most recent call last)
<ipython-input-29-2eaa588b03d9> in <module>()
----> 1 df.B[(df.A > 0) & (df.A < 3)].sum().compute()

/home/ec2-user/anaconda2/lib/python2.7/site-packages/dask/dataframe/core.pyc in sum(self, axis, skipna)
   1369     @derived_from(pd.Series)
   1370     def sum(self, axis=None, skipna=True):
-> 1371         return super(Series, self).sum(axis=axis, skipna=skipna)
   1372 
   1373     @derived_from(pd.Series)

/home/ec2-user/anaconda2/lib/python2.7/site-packages/dask/dataframe/core.pyc in sum(self, axis, skipna)
    866     def sum(self, axis=None, skipna=True):
    867         axis = self._validate_axis(axis)
--> 868         meta = self._meta_nonempty.sum(axis=axis, skipna=skipna)
    869         if axis == 1:
    870             return map_partitions(_sum, self, meta=meta,

TypeError: _sum() got an unexpected keyword argument 'skipna'

Using dask 0.11 on Amazon Linux.

What is your version of Pandas and Dask?

On Fri, Oct 7, 2016 at 3:00 AM, Mauricio Roman notifications@github.com
wrote:

> I am also getting a similar error when working through the example
> Notebooks:
> 
> ## df.B[(df.A > 0) & (df.A < 3)].sum().compute()
> 
> TypeError Traceback (most recent call last)
> in ()
> ----> 1 df.B[(df.A > 0) & (df.A < 3)].sum().compute()
> 
> /home/ec2-user/anaconda2/lib/python2.7/site-packages/dask/dataframe/core.pyc
> in sum(self, axis, skipna)
> 1369 @derived_from(pd.Series)
> 1370 def sum(self, axis=None, skipna=True):
> -> 1371 return super(Series, self).sum(axis=axis, skipna=skipna)
> 1372
> 1373 @derived_from(pd.Series)
> 
> /home/ec2-user/anaconda2/lib/python2.7/site-packages/dask/dataframe/core.pyc
> in sum(self, axis, skipna)
> 866 def sum(self, axis=None, skipna=True):
> 867 axis = self._validate_axis(axis)
> --> 868 meta = self._meta_nonempty.sum(axis=axis, skipna=skipna)
> 869 if axis == 1:
> 870 return map_partitions(_sum, self, meta=meta,
> 
> TypeError: _sum() got an unexpected keyword argument 'skipna'
> 
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> https://github.com/dask/dask/issues/1572#issuecomment-252167616, or mute
> the thread
> https://github.com/notifications/unsubscribe-auth/AASszCms0DEaxL3cv-bWQhMf7X-qO8Duks5qxe4EgaJpZM4KEtGo
> .

@MauricioRoman : I can't seem to replicate on dask 0.11.0, or on the latest release (0.11.1). I suggest reinstalling dask and seeing if it clears up. If not, please open a new issue, this is unrelated to the issue being discussed here (which is about `dask.bag`).


Resolving patch:
diff --git a/dask/bag/__init__.py b/dask/bag/__init__.py
--- a/dask/bag/__init__.py
+++ b/dask/bag/__init__.py
@@ -1,7 +1,8 @@
 from __future__ import absolute_import, division, print_function
 
 from .core import (Bag, Item, from_sequence, from_url, to_textfiles, concat,
-                   from_delayed, bag_range as range, bag_zip as zip)
+                   from_delayed, bag_range as range, bag_zip as zip,
+                   bag_map as map)
 from .text import read_text
 from ..context import set_options
 from ..base import compute
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -34,6 +34,7 @@
 from ..compatibility import apply, urlopen
 from ..context import _globals, defer_to_globals
 from ..core import quote, istask, get_dependencies, reverse_dict
+from ..delayed import Delayed
 from ..multiprocessing import get as mpget
 from ..optimize import fuse, cull, inline, dont_optimize
 from ..utils import (open, system_encoding, takes_multiple_arguments, funcname,
@@ -380,52 +381,126 @@ def __str__(self):
 
     str = property(fget=StringAccessor)
 
-    def map(self, func, **kwargs):
-        """ Map a function across all elements in collection
+    def map(self, func, *args, **kwargs):
+        """Apply a function elementwise across one or more bags.
 
+        Note that all ``Bag`` arguments must be partitioned identically.
+
+        Parameters
+        ----------
+        func : callable
+        *args, **kwargs : Bag, Item, or object
+            Extra arguments and keyword arguments to pass to ``func`` *after*
+            the calling bag instance. Non-Bag args/kwargs are broadcasted
+            across all calls to ``func``.
+
+        Notes
+        -----
+        For calls with multiple `Bag` arguments, corresponding partitions
+        should have the same length; if they do not, the "extra" elements from
+        the longer partition(s) will be dropped.
+
+        Examples
+        --------
         >>> import dask.bag as db
-        >>> b = db.from_sequence(range(5))
-        >>> list(b.map(lambda x: x * 10))  # doctest: +SKIP
-        [0, 10, 20, 30, 40]
+        >>> b = db.from_sequence(range(5), npartitions=2)
+        >>> b2 = db.from_sequence(range(5, 10), npartitions=2)
 
-        Keyword arguments are passed through to ``func``. These can be either
-        ``dask.bag.Item``, or normal python objects.
+        Apply a function to all elements in a bag:
+
+        >>> b.map(lambda x: x + 1).compute()
+        [1, 2, 3, 4, 5]
+
+        Apply a function with arguments from multiple bags:
+
+        >>> from operator import add
+        >>> b.map(add, b2).compute()
+        [5, 7, 9, 11, 13]
+
+        Non-bag arguments are broadcast across all calls to the mapped
+        function:
+
+        >>> b.map(add, 1).compute()
+        [1, 2, 3, 4, 5]
+
+        Keyword arguments are also supported, and have the same semantics as
+        regular arguments:
+
+        >>> def myadd(x, y=0):
+        ...     return x + y
+        >>> b.map(myadd, y=b2).compute()
+        [5, 7, 9, 11, 13]
+        >>> b.map(myadd, y=1).compute()
+        [1, 2, 3, 4, 5]
+
+        Both arguments and keyword arguments can also be instances of
+        ``dask.bag.Item``. Here we'll add the max value in the bag to each
+        element:
+
+        >>> b.map(myadd, b.max()).compute()
+        [4, 5, 6, 7, 8]
+        """
+        if takes_multiple_arguments(func, varargs=False) and not args:
+            warn("Automatic 'splatting' of arguments in `Bag.map` is "
+                 "deprecated and will be removed in a future release. "
+                 "Please use `Bag.starmap` instead.")
+            return self.starmap(func, **kwargs)
+        else:
+            return bag_map(func, self, *args, **kwargs)
+
+    def starmap(self, func, **kwargs):
+        """Apply a function using argument tuples from the given bag.
+
+        This is similar to ``itertools.starmap``, except it also accepts
+        keyword arguments. In pseudocode, this is could be written as:
+
+        >>> def starmap(func, bag, **kwargs):
+        ...     return (func(*args, **kwargs) for args in bag)
+
+        Parameters
+        ----------
+        func : callable
+        **kwargs : Item, Delayed, or object, optional
+            Extra keyword arguments to pass to ``func``. These can either be
+            normal objects, ``dask.bag.Item``, or ``dask.delayed.Delayed``.
 
         Examples
         --------
         >>> import dask.bag as db
-        >>> b = db.from_sequence(range(1, 101), npartitions=10)
-        >>> def div(num, den=1):
-        ...     return num / den
+        >>> data = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
+        >>> b = db.from_sequence(data, npartitions=2)
 
-        Using a python object:
+        Apply a function to each argument tuple:
 
-        >>> hi = b.max().compute()
-        >>> hi
-        100
-        >>> b.map(div, den=hi).take(5)
-        (0.01, 0.02, 0.03, 0.04, 0.05)
+        >>> from operator import add
+        >>> b.starmap(add).compute()
+        [3, 7, 11, 15, 19]
 
-        Using an ``Item``:
+        Apply a function to each argument tuple, with additional keyword
+        arguments:
 
-        >>> b.map(div, den=b.max()).take(5)
-        (0.01, 0.02, 0.03, 0.04, 0.05)
+        >>> def myadd(x, y, z=0):
+        ...     return x + y + z
+        >>> b.starmap(myadd, z=10).compute()
+        [13, 17, 21, 25, 29]
 
-        Note that while both versions give the same output, the second forms a
-        single graph, and then computes everything at once, and in some cases
-        may be more efficient.
+        Keyword arguments can also be instances of ``dask.bag.Item`` or
+        ``dask.delayed.Delayed``:
+
+        >>> max_second = b.pluck(1).max()
+        >>> max_second.compute()
+        10
+        >>> b.starmap(myadd, z=max_second).compute()
+        [13, 17, 21, 25, 29]
         """
-        name = 'map-{0}-{1}'.format(funcname(func),
-                                    tokenize(self, func, kwargs))
-        if takes_multiple_arguments(func):
-            func = partial(apply, func)
+        name = 'starmap-{0}-{1}'.format(funcname(func),
+                                        tokenize(self, func, kwargs))
         dsk = self.dask.copy()
         if kwargs:
-            kw_dsk, kw_pairs = unpack_kwargs(kwargs)
+            kw_dsk, kwargs = unpack_scalar_dask_kwargs(kwargs)
             dsk.update(kw_dsk)
-            func = (apply, partial, [func], (dict, kw_pairs))
-        dsk.update(((name, i), (reify, (map, func, (self.name, i))))
-                   for i in range(self.npartitions))
+        dsk.update({(name, i): (reify, (starmap_chunk, func, (self.name, i), kwargs))
+                   for i in range(self.npartitions)})
         return type(self)(dsk, name, self.npartitions)
 
     @property
@@ -513,7 +588,7 @@ def map_partitions(self, func, **kwargs):
         >>> b.map_partitions(myfunc)  # doctest: +SKIP
 
         Keyword arguments are passed through to ``func``. These can be either
-        ``dask.bag.Item``, or normal python objects.
+        ``dask.bag.Item``, ``dask.delayed.Delayed``, or normal python objects.
 
         Examples
         --------
@@ -543,12 +618,12 @@ def map_partitions(self, func, **kwargs):
                                                tokenize(self, func, kwargs))
         dsk = self.dask.copy()
         if kwargs:
-            kw_dsk, kw_pairs = unpack_kwargs(kwargs)
+            kw_dsk, kwargs = unpack_scalar_dask_kwargs(kwargs)
             dsk.update(kw_dsk)
-        dsk.update(((name, i),
-                    (apply, func, [(self.name, i)], (dict, kw_pairs))
-                    if kwargs else (func, (self.name, i)))
-                   for i in range(self.npartitions))
+            task = lambda i: (apply, func, [(self.name, i)], kwargs)
+        else:
+            task = lambda i: (func, (self.name, i))
+        dsk.update(((name, i), task(i)) for i in range(self.npartitions))
         return type(self)(dsk, name, self.npartitions)
 
     def pluck(self, key, default=no_default):
@@ -1507,6 +1582,186 @@ def bag_zip(*bags):
     return Bag(merge(bags_dsk, dsk), name, npartitions)
 
 
+def map_chunk(f, args, bag_kwargs, kwargs):
+    if kwargs:
+        f = partial(f, **kwargs)
+
+    args = [iter(a) for a in args]
+    iters = list(args)
+    if bag_kwargs:
+        keys = list(bag_kwargs)
+        kw_val_iters = [iter(v) for v in bag_kwargs.values()]
+        iters.extend(kw_val_iters)
+        kw_iter = (dict(zip(keys, k)) for k in zip(*kw_val_iters))
+        if args:
+            for a, k in zip(zip(*args), kw_iter):
+                yield f(*a, **k)
+        else:
+            for k in kw_iter:
+                yield f(**k)
+    else:
+        for a in zip(*args):
+            yield f(*a)
+
+    # Check that all iterators are fully exhausted
+    if len(iters) > 1:
+        for i in iters:
+            if isinstance(i, itertools.repeat):
+                continue
+            try:
+                next(i)
+            except StopIteration:
+                pass
+            else:
+                msg = ("map called with multiple bags that aren't identically "
+                       "partitioned. Please ensure that all bag arguments "
+                       "have the same partition lengths")
+                raise ValueError(msg)
+
+
+def starmap_chunk(f, x, kwargs):
+    if kwargs:
+        f = partial(f, **kwargs)
+    return itertools.starmap(f, x)
+
+
+def unpack_scalar_dask_kwargs(kwargs):
+    """Extracts dask values from kwargs
+
+    Currently only ``dask.bag.Item`` and ``dask.delayed.Delayed`` are
+    supported.  Returns a merged dask graph and a task resulting in a keyword
+    dict.
+    """
+    dsk = {}
+    kwargs2 = {}
+    for k, v in kwargs.items():
+        if isinstance(v, (Delayed, Item)):
+            dsk.update(v.dask)
+            kwargs2[k] = v.key
+        elif isinstance(v, Base):
+            raise NotImplementedError("dask.bag doesn't support kwargs of "
+                                      "type %s" % type(v).__name__)
+        else:
+            kwargs2[k] = v
+    if dsk:
+        kwargs = (dict, (zip, list(kwargs2), list(kwargs2.values())))
+    return dsk, kwargs
+
+
+def bag_map(func, *args, **kwargs):
+    """Apply a function elementwise across one or more bags.
+
+    Note that all ``Bag`` arguments must be partitioned identically.
+
+    Parameters
+    ----------
+    func : callable
+    *args, **kwargs : Bag, Item, Delayed, or object
+        Arguments and keyword arguments to pass to ``func``. Non-Bag args/kwargs
+        are broadcasted across all calls to ``func``.
+
+    Notes
+    -----
+    For calls with multiple `Bag` arguments, corresponding partitions should
+    have the same length; if they do not, the "extra" elements from the longer
+    partition(s) will be dropped.
+
+    Examples
+    --------
+    >>> import dask.bag as db
+    >>> b = db.from_sequence(range(5), npartitions=2)
+    >>> b2 = db.from_sequence(range(5, 10), npartitions=2)
+
+    Apply a function to all elements in a bag:
+
+    >>> db.map(lambda x: x + 1, b).compute()
+    [1, 2, 3, 4, 5]
+
+    Apply a function with arguments from multiple bags:
+
+    >>> from operator import add
+    >>> db.map(add, b, b2).compute()
+    [5, 7, 9, 11, 13]
+
+    Non-bag arguments are broadcast across all calls to the mapped function:
+
+    >>> db.map(add, b, 1).compute()
+    [1, 2, 3, 4, 5]
+
+    Keyword arguments are also supported, and have the same semantics as
+    regular arguments:
+
+    >>> def myadd(x, y=0):
+    ...     return x + y
+    >>> db.map(myadd, b, y=b2).compute()
+    [5, 7, 9, 11, 13]
+    >>> db.map(myadd, b, y=1).compute()
+    [1, 2, 3, 4, 5]
+
+    Both arguments and keyword arguments can also be instances of
+    ``dask.bag.Item`` or ``dask.delayed.Delayed``. Here we'll add the max value
+    in the bag to each element:
+
+    >>> db.map(myadd, b, b.max()).compute()
+    [4, 5, 6, 7, 8]
+    """
+    name = 'map-%s-%s' % (funcname(func), tokenize(func, args, kwargs))
+    dsk = {}
+
+    bags = []
+    args2 = []
+    for a in args:
+        if isinstance(a, Bag):
+            bags.append(a)
+            args2.append(a)
+            dsk.update(a.dask)
+        elif isinstance(a, (Item, Delayed)):
+            args2.append((itertools.repeat, a.key))
+            dsk.update(a.dask)
+        else:
+            args2.append((itertools.repeat, a))
+
+    bag_kwargs = {}
+    other_kwargs = {}
+    for k, v in kwargs.items():
+        if isinstance(v, Bag):
+            bag_kwargs[k] = v
+            bags.append(v)
+            dsk.update(v.dask)
+        else:
+            other_kwargs[k] = v
+
+    kw_dsk, other_kwargs = unpack_scalar_dask_kwargs(other_kwargs)
+    dsk.update(kw_dsk)
+
+    if not bags:
+        raise ValueError("At least one argument must be a Bag.")
+
+    npartitions = {b.npartitions for b in bags}
+    if len(npartitions) > 1:
+        raise ValueError("All bags must have the same number of partitions.")
+    npartitions = npartitions.pop()
+
+    def build_args(n):
+        return [(a.name, n) if isinstance(a, Bag) else a for a in args2]
+
+    def build_bag_kwargs(n):
+        if not bag_kwargs:
+            return None
+        return (dict, (zip, list(bag_kwargs),
+                       [(b.name, n) for b in bag_kwargs.values()]))
+
+    dsk.update({(name, n): (reify, (map_chunk, func, build_args(n),
+                                    build_bag_kwargs(n), other_kwargs))
+                for n in range(npartitions)})
+
+    # If all bags are the same type, use that type, otherwise fallback to Bag
+    return_type = set(map(type, bags))
+    return_type = return_type.pop() if len(return_type) == 1 else Bag
+
+    return return_type(dsk, name, npartitions)
+
+
 def _reduce(binop, sequence, initial=no_default):
     if initial is not no_default:
         return reduce(binop, sequence, initial)
diff --git a/dask/utils.py b/dask/utils.py
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -494,7 +494,7 @@ def file_size(fn, compression=None):
                             isinstance, issubclass, map, pow, setattr])
 
 
-def takes_multiple_arguments(func):
+def takes_multiple_arguments(func, varargs=True):
     """ Does this function take multiple arguments?
 
     >>> def f(x, y): pass
@@ -534,12 +534,11 @@ def takes_multiple_arguments(func):
     except:
         is_constructor = False
 
-    if spec.varargs:
+    if varargs and spec.varargs:
         return True
 
-    if spec.defaults is None:
-        return len(spec.args) - is_constructor != 1
-    return len(spec.args) - len(spec.defaults) - is_constructor > 1
+    ndefaults = 0 if spec.defaults is None else len(spec.defaults)
+    return len(spec.args) - ndefaults - is_constructor > 1
 
 
 class Dispatch(object):


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -47,14 +47,14 @@ def test_keys():
     assert sorted(b._keys()) == sorted(dsk.keys())
 
 
+# Deprecated test for old map behavior
 def test_map():
     c = b.map(inc)
-    expected = merge(dsk, dict(((c.name, i), (reify, (map, inc, (b.name, i))))
-                               for i in range(b.npartitions)))
-    assert c.dask == expected
+    assert c.compute() == list(map(inc, b.compute()))
     assert c.name == b.map(inc).name
 
 
+# Deprecated test for old map behavior
 def test_map_function_with_multiple_arguments():
     b = db.from_sequence([(1, 10), (2, 20), (3, 30)], npartitions=3)
     assert list(b.map(lambda x, y: x + y).compute(get=dask.get)) == [11, 22, 33]
@@ -71,6 +71,7 @@ def __init__(self, a):
         pass
 
 
+# Deprecated test for old map behavior
 def test_map_with_constructors():
     assert db.from_sequence([[1, 2, 3]]).map(A).compute()
     assert db.from_sequence([1, 2, 3]).map(B).compute()
@@ -84,6 +85,7 @@ def test_map_with_constructors():
     assert failed
 
 
+# Deprecated test for old map behavior
 def test_map_with_builtins():
     b = db.from_sequence(range(3))
     assert ' '.join(b.map(str)) == '0 1 2'
@@ -98,6 +100,7 @@ def test_map_with_builtins():
         [True, False]
 
 
+# Deprecated test for old map behavior
 def test_map_with_kwargs():
     b = db.from_sequence(range(100), npartitions=10)
     assert b.map(lambda x, factor=0: x * factor,
@@ -109,6 +112,105 @@ def test_map_with_kwargs():
                  factor=2).sum().compute() == 2.0
 
 
+def test_bag_map():
+    b = db.from_sequence(range(100), npartitions=10)
+    b2 = db.from_sequence(range(100, 200), npartitions=10)
+    x = b.compute()
+    x2 = b2.compute()
+
+    def myadd(a=1, b=2, c=3):
+        return a + b + c
+
+    assert db.map(myadd, b).compute() == list(map(myadd, x))
+    assert db.map(myadd, a=b).compute() == list(map(myadd, x))
+    assert db.map(myadd, b, b2).compute() == list(map(myadd, x, x2))
+    assert db.map(myadd, b, 10).compute() == [myadd(i, 10) for i in x]
+    assert db.map(myadd, 10, b=b).compute() == [myadd(10, b=i) for i in x]
+
+    sol = [myadd(i, b=j, c=100) for (i, j) in zip(x, x2)]
+    assert db.map(myadd, b, b=b2, c=100).compute() == sol
+
+    sol = [myadd(i, c=100) for (i, j) in zip(x, x2)]
+    assert db.map(myadd, b, c=100).compute() == sol
+
+    x_sum = sum(x)
+    sol = [myadd(x_sum, b=i, c=100) for i in x2]
+    assert db.map(myadd, b.sum(), b=b2, c=100).compute() == sol
+
+    sol = [myadd(i, b=x_sum, c=100) for i in x2]
+    assert db.map(myadd, b2, b.sum(), c=100).compute() == sol
+
+    sol = [myadd(a=100, b=x_sum, c=i) for i in x2]
+    assert db.map(myadd, a=100, b=b.sum(), c=b2).compute() == sol
+
+    a = dask.delayed(10)
+    assert db.map(myadd, b, a).compute() == [myadd(i, 10) for i in x]
+    assert db.map(myadd, b, b=a).compute() == [myadd(i, b=10) for i in x]
+
+    # Mispatched npartitions
+    fewer_parts = db.from_sequence(range(100), npartitions=5)
+    with pytest.raises(ValueError):
+        db.map(myadd, b, fewer_parts)
+
+    # No bags
+    with pytest.raises(ValueError):
+        db.map(myadd, b.sum(), 1, 2)
+
+    # Unequal partitioning
+    unequal = db.from_sequence(range(110), npartitions=10)
+    with pytest.raises(ValueError):
+        db.map(myadd, b, unequal, c=b2).compute()
+    with pytest.raises(ValueError):
+        db.map(myadd, b, b=unequal, c=b2).compute()
+
+
+def test_map_method():
+    b = db.from_sequence(range(100), npartitions=10)
+    b2 = db.from_sequence(range(100, 200), npartitions=10)
+    x = b.compute()
+    x2 = b2.compute()
+
+    def myadd(a, b=2, c=3):
+        return a + b + c
+
+    assert b.map(myadd).compute() == list(map(myadd, x))
+    assert b.map(myadd, b2).compute() == list(map(myadd, x, x2))
+    assert b.map(myadd, 10).compute() == [myadd(i, 10) for i in x]
+    assert b.map(myadd, b=10).compute() == [myadd(i, b=10) for i in x]
+    assert (b.map(myadd, b2, c=10).compute() ==
+            [myadd(i, j, 10) for (i, j) in zip(x, x2)])
+    x_sum = sum(x)
+    assert (b.map(myadd, b.sum(), c=10).compute() ==
+            [myadd(i, x_sum, 10) for i in x])
+
+    # check that map works with multiarg functions. Can be removed after
+    # deprecated behavior is removed
+    assert b.map(add, b2).compute() == list(map(add, x, x2))
+
+    # check that map works with vararg functions. Can be removed after
+    # deprecated behavior is removed
+    def vararg_inc(*args):
+        return inc(*args)
+
+    assert b.map(vararg_inc).compute(get=dask.get) == list(map(inc, x))
+
+
+def test_starmap():
+    data = [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10)]
+    b = db.from_sequence(data, npartitions=2)
+
+    def myadd(a, b, c=0):
+        return a + b + c
+
+    assert b.starmap(myadd).compute() == [myadd(*a) for a in data]
+    assert b.starmap(myadd, c=10).compute() == [myadd(*a, c=10) for a in data]
+    max_second = b.pluck(1).max()
+    assert (b.starmap(myadd, c=max_second).compute() ==
+            [myadd(*a, c=max_second.compute()) for a in data])
+    c = dask.delayed(10)
+    assert b.starmap(myadd, c=c).compute() == [myadd(*a, c=10) for a in data]
+
+
 def test_filter():
     c = b.filter(iseven)
     expected = merge(dsk, dict(((c.name, i),
@@ -334,16 +436,15 @@ def test_map_partitions():
 
 def test_map_partitions_with_kwargs():
     b = db.from_sequence(range(100), npartitions=10)
-    assert b.map_partitions(
-        lambda X, factor=0: [x * factor for x in X],
-        factor=2).sum().compute() == 9900.0
-    assert b.map_partitions(
-        lambda X, total=0: [x / total for x in X],
-        total=b.sum()).sum().compute() == 1.0
-    assert b.map_partitions(
-        lambda X, factor=0, total=0: [x * factor / total for x in X],
-        total=b.sum(),
-        factor=2).sum().compute() == 2.0
+
+    def scale(X, factor=1, total=1):
+        return [x * factor / total for x in X]
+
+    assert b.map_partitions(scale, factor=2).sum().compute() == 9900.0
+    assert b.map_partitions(scale, total=b.sum()).sum().compute() == 1.0
+    assert b.map_partitions(scale, total=b.sum(), factor=2).sum().compute() == 2.0
+    total = dask.delayed(sum)(list(range(100)))
+    assert b.map_partitions(scale, total=total, factor=2).sum().compute() == 2.0
 
 
 def test_random_sample_size():
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -68,6 +68,17 @@ def __init__(self, a, b):
     assert not takes_multiple_arguments(Singular)
     assert takes_multiple_arguments(Multi)
 
+    def f():
+        pass
+
+    assert not takes_multiple_arguments(f)
+
+    def vararg(*args):
+        pass
+
+    assert takes_multiple_arguments(vararg)
+    assert not takes_multiple_arguments(vararg, varargs=False)
+
 
 def test_dispatch():
     foo = Dispatch()