forked from inaos/iron-array-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathudf_expr.py
More file actions
66 lines (52 loc) · 1.68 KB
/
Copy pathudf_expr.py
File metadata and controls
66 lines (52 loc) · 1.68 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
# Calling scalar UDFs from expressions.
# This is for 1-dim arrays.
from time import time
import numpy as np
import iarray as ia
from iarray import udf
# Define array params
shape = [100_000_000]
@udf.scalar(lib="lib")
def fsum(a: udf.float64, b: udf.float64) -> float:
return a + b
@udf.scalar(lib="lib2")
def fmult(a: udf.float64, b: udf.float64) -> float:
return a * b
# With the new mechanism for registering functions in the udf.scalar decorator,
# it is not necessary to register functions manually anymore.
# ia.udf_registry["lib"] = fsum
# ia.udf_registry["lib2"] = fmult
print("Registered UDF funcs:", tuple(ia.udf_registry.iter_all_func_names()))
# Create initial containers
a1 = ia.linspace(shape, 0, 10)
a2 = np.linspace(0, 10, shape[0]).reshape(shape)
print("** pure expr evaluation ...")
expr = "4 * (x * y)"
expr = ia.expr_from_string(expr, {"x": a1, "y": 1})
t0 = time()
b1 = expr.eval()
print("Time:", round((time() - t0), 3))
print(f"cratio for result: {b1.cratio:.3f}")
b1_n = b1.data
print(b1_n)
print("** scalar udf evaluation ...")
# expr = "lib.f(a1, a1)" # segfault. fix it by propagating errors correctly!
# expr = "4 * lib.fsum(x, x) + lib2.fmult(x, x)" # segfaults too
expr = "4 * lib2.fmult(x, y)"
expr = ia.expr_from_string(expr, {"x": a1, "y": 1})
t0 = time()
b1 = expr.eval()
print("Time:", round((time() - t0), 3))
print(f"cratio for result: {b1.cratio:.3f}")
b1_n = b1.data
print(b1_n)
import numexpr as ne
print("** numexpr evaluation ...")
expr = "4 * x * y"
a1_n = a1.data
t0 = time()
b1 = ne.evaluate(expr, {"x": a1_n, "y": 1})
print("Time:", round((time() - t0), 3))
print(b1)
# In case we want to clear the UDF registry explicitly
ia.udf_registry.clear()