Computing expressions like numexpr.evaluate or pandas.eval

[kreczko]

Version of Awkward Array

2.0.7

Description and code to reproduce

I am trying to use pandas.eval (pandas == 1.5.3) with awkward arrays.

The minimal example for this is

import awkward as ak
import pandas as pd
expression = "a + b"
a = ak.Array([[1, 2, 3], [5, 6]])
b = ak.Array([[2, 3, 4], [6, 7]])
result = pd.eval(expression, engine="numexpr")
print(result)

The error I am getting is:

<snip>
  File "/software/miniconda/envs/fast-iris/lib/python3.10/site-packages/pandas/core/computation/ops.py", line 166, in is_datetime
    return issubclass(t, (datetime, np.datetime64))
TypeError: issubclass() arg 1 must be a class

Patching the pandas code reveals that the value t in is_datetime is

<property object at 0x7f5399ec4860>

Replacing awkward arrays with numpy arrays works (t = <class 'numpy.int64'>).

Not sure of the exact chain of events, but before I dig deeper, I thought I post this in case someone has an "of course, it's that thing" moment ;).
If you need me to drill deeper, please let me know what you need.

PS: Same issue with awkward 1.X, so I guess nobody tried pandas.eval with awkward arrays so far.

Note: The reason I am trying to use pandas.eval over numexpr or awkward..evaluate is the ease of using custom functions (with all its downsides).

[jpivarski]

I suppose that pd.eval is running numexpr.evaluate, which grabs variables by name from the surrounding environment, assumes that they are NumPy arrays, and performs the calculation, returning a NumPy array.

Awkward v1 had an ak.numexpr.evaluate, which I wasn't very happy with because it's in the ak.* namespace instead of numexpr.evaluate just knowing what to do with Awkward Arrays. So although it has been updated and is hidden in the v2 code, it hasn't been exposed publicly yet.

awkward/src/awkward/_connect/numexpr.py

Lines 69 to 71 in 6b49960

	def evaluate(
	expression, local_dict=None, global_dict=None, order="K", casting="safe", **kwargs
	):

Before just making this public, I wanted to see if there's any way to register Awkward with NumExpr (like the way that it's registered with NumPy's ufuncs or Numba's JIT). I think that was something that was going to be added to NumExpr v3, but it doesn't look like NumExpr v3 is going to happen.

NumExpr expressions would be great to turn into ufuncs, and then they could be used uniformly on any array type that extends NumPy (NEP-13, __array_ufunc__). Let me think about that...

[kreczko]

Hi Jim,

suppose that pd.eval is running numexpr.evaluate, > which grabs variables by name from the surrounding environment, assumes that they are NumPy arrays, and performs the calculation, returning a NumPy array.

Yes, I think that's what it does (with probably slight variations for different engines).

has been updated and is hidden in the v2 code

Not well enough, tried this one yesterday ;).

The main issue for me are custom functions, e.g.

def func2(a, b):
    return a + b

expression = "a + b + func2(a, b)"

which works for pandas.eval but cannot be done with ak._connect.numexpr.evaluate nor numexpr.evaluate (although in most cases things like func2 would point to ak functions).

Anyway, not part of this issue ;).

Let me think about that...

Of course :). Let me know if you need a hand or two. Beats writing a custom parser for these expressions.

[jpivarski]

Not well enough, tried this one yesterday

No, I mean it's still hidden; you can't use it (without going through a module name that starts with an underscore).

The one thing we can't do is make this slick, to have ne.evaluate or pd.eval just pick up the Awkward Arrays and know what to do with them.

I'm also assuming that this isn't what you're looking for:

>>> a = ak.Array([[1, 2, 3], [5, 6]])
>>> b = ak.Array([[2, 3, 4], [6, 7]])
>>> eval("a + b")
<Array [[3, 5, 7], [11, 13]] type='2 * var * int64'>

Is it about wanting to do a single pass over the data? With something like the above (execution of the string in Python), a + b + c would first iterate (in precompiled code) over a and b, adding each, and then iterate over the sum and c, adding each. The reason for using NumExpr is to avoid all of the intermediate NumPy arrays and separate passes over the data (which uses CPU caches badly).

If you're okay with an interface like this,

>>> evaluate("a + b", a=a, b=b)
<Array [[3, 5, 7], [11, 13]] type='2 * var * int64'>

it can be done like this:

def evaluate(expression, **namespace):
    arrays = {}
    not_arrays = {}
    for k, v in namespace.items():
        if isinstance(v, (ak.Array, np.array)):  # maybe also pd.Series...
            arrays[k] = v
        else:
            not_arrays[k] = v

    def action(layouts, **ignore):
        if all(x.is_numpy for x in layouts):
            ns = dict(not_arrays)
            for name, value in zip(arrays.keys(), layouts):
                ns[name] = value.data
            return ak.contents.NumpyArray(ne.evaluate(expression, ns))

    return ak.transform(action, *arrays.values())

This is, in a sense, a reimplementation of ak.numexpr.evaluate using ak.transform. It applies all of the broadcasting rules to a set of arrays (arrays.values()) while recursing down until you get to the level of ak.contents.NumpyArray, which wrap plain NumPy arrays. Then the action function determines what to do with those flat arrays: in this case, it fills a namespace (ns), calls ne.evaluate (can be pd.eval) and wraps the result as a new ak.contents.NumpyArray that is part of the output.

If what you really want is to get the pd.eval function to work, then I don't see how to do that because we can't override Pandas.

[jpivarski]

If the part that you really want is to pick up the variable names from the string, that is possible if it can be parsed. ast.parse would do it if the language (NumExpr) is a subset of the Python language, which I think it is.

>>> print(ast.dump(ast.parse("a + b"), indent="    "))
Module(
    body=[
        Expr(
            value=BinOp(
                left=Name(id='a', ctx=Load()),
                op=Add(),
                right=Name(id='b', ctx=Load())))],
    type_ignores=[])

One could then write a Visitor to walk over that AST and find the Name objects that are not in the function part of a Call (because you want non-arrays, like func2, to not be interpreted as arrays).

[kreczko]

Thanks a lot Jim!

without going through a module name that starts with an underscore

Which is fine for single underscores:

result = ak._connect.numexpr.evaluate(expression)

→ works (for expressions without functions)

Is it about wanting to do a single pass over the data?

Yes. Although I care at the moment more about completeness than efficiency.

Regarding your suggestions, indeed, these are the paths I tried over the weekend - and then I remembered pandas.eval ;)

Anyway, you are right, a custom implementation would be best as I wanted to extract a compute graph in its final form.
Forgot about ast.parse was trying out ply instead :D.

For ast it does indeed give me something very useful for print(ast.dump(ast.parse("a + b + func2(a, b)"), indent=" ")):

Module(
  body=[
    Expr(
      value=BinOp(
        left=BinOp(
          left=Name(id='a', ctx=Load()),
          op=Add(),
          right=Name(id='b', ctx=Load())),
        op=Add(),
        right=Call(
          func=Name(id='func2', ctx=Load()),
          args=[
            Name(id='a', ctx=Load()),
            Name(id='b', ctx=Load())],
          keywords=[])))],
  type_ignores=[])

I can use this information to create a compute graph or extract & evaluate the functions to make the expression compatible with numexpr.

Thanks for the pointers!