How to make custom object available for function passed to dask df.apply (cannot serialize)

user_1415130 0
Viewed 1148

All this code works in pandas, but running single threaded is slow.

I have an object (it's a bloom filter) that's slow to create.

I have dask code that looks something like:

def has_match(row, my_filter):
    return my_filter.matches(
        a=row.a, b =row.b
    )

# ....make dask dataframe ddf

ddf['match'] = ddf.apply(has_match, args=(my_filter, ), axis=1, meta=(bool))
ddf.compute()

When I try to run this I get an error that starts:

distributed.protocol.core - CRITICAL - Failed to Serialize

My object was created from a C library, so I'm not surprised that it can't be automagically serialized, but I don't know how to work around this.

dask-distributed dask
2 Answers

Use only threads

One way is to avoid the problem altogether and just not use separate processes at all. That way you won't need to serialize data between them.

ddf.compute(scheduler='threads')

This does limit you to running in a single process on a single machine though, which may not be what you want.

Figure out how to serialize your object

If you can figure out how to turn your object into a bytestring and back then you can either implement the pickle protocol on your object (like the __getstate__ and __setstate__ methods, see Python docs) or you can add definitions to the dask_serialize and dask_deserialize dispatchable functions. See Dask's serialization docs for an example.

Re-create your objects every time

Maybe it's hard to serialize your object, but cheap to recreate it once per partition?

def has_match(partition):
    my_filter = make_filter(...)
    return partition.apply(my_filter.matches(a=row.a, b =row.b))

ddf['match'] = ddf.map_partitions(has_match)
user_6166160

Distributed expects all intermediate results to be serializable. In your case, you have an object that doesn't implement pickle. In general you have a few options here (in order of best to worst IMHO):

  • Implement pickle for this object. Note that using the copyreg module you can add pickle support for classes that aren't in your control.

  • Cache the creation of the filter in your function manually. You could do this with an object, or with a global variable in your module. Note that this code below would need to be part of a module imported, not part of your interactive session (i.e. not in a jupyter notebook/ipython session).

For example (untested):

myfilter = None


def get_or_load():
    global myfilter
    if myfilter is None:
        myfilter = load_filter()
    else:
        return myfilter


def load_filter():
    pass


def has_match(row):
    my_filter = get_or_load()
    return my_filter.matches(a=row.a, b=row.b)

And then in your user code:

from my_filter_utils import has_match

ddf['match'] = ddf.apply(has_match, axis=1, meta=('matches', bool))
  • Use dask to manage the cache. To do this, wrap the object in another class that re-loads the object when serialized. If you then persist that object in the cluster, dask will hold on to it and at most the creation function will be called once on every node.

For example (untested):

from dask import delayed

class Wrapper(object):
    def __init__(self, func):
        self.func = func
        self.filter = func()

    def __reduce__(self):
        # When unpickled, the filter will be reloaded
        return (Wrapper, (func,))


def load_filter():
    pass


# Create a delayed function to load the filter
wrapper = delayed(Wrapper)(load_filter)

# Optionally persist the wrapper in the cluster, to be reused over multiple computations
wrapper = wrapper.persist()

def has_match(row, wrapper):
    return wrapper.filter.matches(a=row.a, b=row.b)


ddf['match'] = ddf.apply(has_match, args=(wrapper,), axis=1, meta=('matches', bool))
user_16672870
Related