weld-project / weld

High-performance runtime for data analytics applications
https://www.weld.rs
BSD 3-Clause "New" or "Revised" License
2.99k stars 258 forks source link

Add function decorator to make Python Weld codegen easier #504

Closed sppalkia closed 4 years ago

sppalkia commented 4 years ago

This patch adds a new function decorator weldfunc that converts a function that returns strings of Weld code into functions that return WeldLazy. The function automatically detects inputs that are themselves WeldLazy, and adds them as dependencies to the return value of the wrapped function. For example, say we have:

def exp(a, b):
   return "exp({}, {})".format(a, b)

This function (or its caller, as is done right now in Grizzly) would ordinarily need to check whether a or b are WeldLazy objects, extract the name, construct the dependency list, etc.:

val1 = 5.0
val2 = ... # create a WeldLazy
code = exp(val1, val2)
dependencies = [val2]
result = WeldLazy(code, dependencies, output_type, decoder)

Instead of doing this manually, decorating exp with weldfunc will construct the dependency list automatically and return a WeldLazy:

@weldfunc # With this, exp can assume all inputs are strings now
def exp(a, b):
   return "exp({}, {})".format(a, b)

# Code to create result as above
val1 = 5.0
val2 = # create a WeldLazy
result = exp(val1, val2)(output_type, decoder)

Note that the return value of exp(val1, val2) is itself a callable, so you can still specify a custom decoder/output type (eventually, the hope is that the output type can be inferred in all cases with a nicer API, but since we have strings for now, this will do). result will already assign val2 as a dependency above. Users can still get the raw code constructed by exp with exp(val1, val2).code (this is mostly just useful for tests).

How is this patch tested?

deepakn94 commented 4 years ago

Cool, feel free to merge!

sppalkia commented 4 years ago

@deepakn94 thanks for the review!