Apply in batches?

[NickCrews]

I have a lot of strings that I want to tokenize for some NLP. I want to use spacy, which has an API like Tokenizer.pipe(Iterable[str]) -> Generator[Doc, None, None]. So I could do documents = Tokenizer.pipe(df["text"].values), but then if I try to materialize the resulting documents into a pyarrow array with pa.array(documents) then I'll run out of memory.

Is there something like DataFrame.apply_batched(func, args) that could operate in batches and build up the result out of core? Or I guess really this might be more like building an arrow array from a generator, out of core, which might be more int he wheelhouse of pyarrow, and not vaex?

Thank you!

[JovanVeljanoski]

Hi,

apply in vaex runs out of core. Vaex takes care of the batches (we call them chunks here).

[NickCrews]

Ahh, I was looking for the vectorize=True option in apply, which I didn't see before. Then it calls f with chunks. I assumed that apply only worked on individual rows/elements, which was not what I was looking for. This is awesome that you support all three modes:

• operating on entire column at once
• If that doesn't work, then operate on chunks using apply(vectorize=True)
• If that doesn't work, operate on individual rows using apply(vectorize=False)

I may submit a PR that improves the docstring of apply that makes this more clear. Thank you!

[JovanVeljanoski]

No no..

apply always works in chunks, i.e. it will always be out-of-core. Everything that you do in vaex is out of core. You never operate on the whole column at once (unless you have a dataframe of length 1, or do something super custom).

apply is a black box - you can put anything there. Certain functions can be made faster by doing vector like operations, e.g.

def my_func(a, b):
    return  (a * b) / (a + b)

For the above case you don't even need apply, but if you were to use it, with vectorize=True the result will be evaluated much faster than with vectorize=False, at least for larger dataframes.

But in other cases, you can't do that, where you can't vectorize operations between elements, so you evaluate the functions on a row-by-row basis. Perhaps @maartenbreddels can explain this much better.

So if you use apply with spacy, you probably want vectorize=False (unless under the hood the process is doing only some linear algebra type of thing, in which case you can give vectorize=True a shot).

[NickCrews]

Thanks for following back up. I think I had that understanding correct, I just phrased it unclearly. By "operating on entire column at once" I meant when you don't use apply, eg df["z"] = (df.a + df.b) / (df.a + df.b). Of course under the hood vaex is still actually performing this computation in chunks, but from an API standpoint that is hidden, it looks like we are operating on the whole column at once.

Actually I think I do want vectorize=True, spacy has two different APIs, one that takes individual strings and processes them, but another that works on iterables of strings and does them in batches. The batch API is faster, so I wanted to use that. This is working for me now:

_nlp = None


def _load_spacy():
    global _nlp
    if _nlp is None:
        _nlp = spacy.load("en_core_web_sm", exclude=["ner", "entity_linker", "textcat"])
    return _nlp


def spacify(s: pa.StringArray) -> pa.ListArray:
    strs = (str(s) for s in s)
    nlp = _load_spacy()
    docs = nlp.pipe(strs)
    token_lists = ([str(tok) for tok in doc] for doc in docs)
    dtype = pa.list_(pa.string())
    return pa.array(token_lists, type=dtype)


df = vaex.from_dict(
    {"text": ["jane's   house", "Ross' house  ", "a", "", None, "bees\tall cook"]}
)
# spacy already does multiprocessing so using multiprocessing doesn't save us anything
df["tokens"] = df.apply(spacify, [df.text], vectorize=True, multiprocessing=False)
df

Gives:

	text	tokens
0	jane's house	['jane' "'s" ' ' 'house']
1	Ross' house	['Ross' "'" 'house' ' ']
2	a	['a']
3		[]
4		['None']
5	bees all cook	['bees' '\t' 'all' 'cook']

One thing that might be nice would be if I could control the chunk_size as an argument here, instead of relying on the global settings, but this isn't a big deal.