Skip to main content
Version: Python

How to use Numba in Python queries

This guide will show you how to use Numba in your Python queries in Deephaven.

Numba is an open-source just-in-time (JIT) compiler for Python. It can be used to translate portions of Python code into optimized machine code using LLVM. The use of Numba can make your queries faster and more responsive.

What is just-in-time (JIT) compilation?​

JIT compiles code into optimized machine code at runtime. When using Numba, you can specify which functions and blocks of code you want to be compiled at runtime. Compiling code blocks into optimized machine code adds some overhead, but subsequent function calls can be much faster. Thus, JIT can be powerful when used on functions that are complex, large, or will be used many times.

caution

JIT does not guarantee faster code. Sometimes, code cannot be optimized well by Numba and may actually run slower.

Usage​

Numba decorators specify which blocks of code to JIT compile. Common Numba decorators include @jit and @vectorize. See http://numba.pydata.org/ for more details.

In the following example, @jit and @vectorize are used to JIT compile functionOne and functionTwo.

import math
from numba import jit, vectorize, double, int64

@jit
def functionOne(a, b):
    return math.sqrt(a**2 + b**2)

@vectorize([double(double, int64)])
def functionTwo(decimal_value, integer_value):
    return decimal_value * integer_value

Lazy compilation​

If the Numba decorator is used without a function signature, Numba will infer the argument types at call time and generate optimized code based upon the inferred types. Numba will also compile separate specializations for different input types. The compilation is deferred to the first function call. This is called lazy compilation.

The example below uses lazy compilation on the function func. When the code is first run, the function is compiled to optimized machine code, which takes extra overhead and results in longer execution time. On the second function call, Python uses the already optimized machine code, which results in very fast execution.

import numpy as np
from numba import jit
import time

# Lazy optimized addition function
@jit
def func(x, y):
    return x + y

x = 0.01 * np.arange(100000)
y = 0.02 * np.arange(100000)

# This function call is slow.  It has to be compiled into optimized machine code.
start = time.time()
z1 = func(x, y)
end = time.time()
print("First function call took " + str(end - start) + " seconds.")

# This function call is fast.  It uses the previously JIT-compiled machine code.
start = time.time()
z2 = func(x, y)
end = time.time()
print("Second function call took " + str(end - start) + " seconds.")

Eager compilation​

If the Numba decorator is used with a function signature, Numba will compile the function when the function is defined. This is called eager compilation.

The example below uses eager compilation on the function func by specifying one or more function signatures. These function signatures denote the allowed input and output data types to and from the function. With eager compilation, the code is compiled into optimized machine code at the function definition. When the code is run the first time, it has already been compiled into optimized machine code, so it's just as fast as the second function call. If the compiled function is used with a data type not specified in the function signature, an error will occur.

from numba import vectorize, int64, double
import numpy as np
import time

# Eager compilation for the int64(int64, int64) and double(double, double) function signatures happens here
@vectorize([int64(int64, int64),
            double(double, double)])
def func(x, y):
    z = x + y
    return z

x1 = 0.01 * np.arange(10000)
y1 = 0.02 * np.arange(10000)
x2 = 1 * np.arange(10000)
y2 = 2 * np.arange(10000)

# Run it the first time on doubles
start = time.time()
z1 = func(x1, y1)
end = time.time()
print("Add vectors of doubles (first run): " + str(end - start) + " seconds.")

# Run it a second time on doubles
start = time.time()
z2 = func(x1, y1)
end = time.time()
print("Add vectors of doubles (second run): " + str(end - start) + " seconds.")

# Run it the first time on integers
start = time.time()
z3 = func(x2, y2)
end = time.time()
print("Add vectors of integers (first run): " + str(end - start) + " seconds.")

# Run it a second time on integers
start = time.time()
z4 = func(x2, y2)
end = time.time()
print("Add vectors of integers (second run): " + str(end - start) + " seconds.")
caution

Eager compilation creates functions that only support specific function signatures. If these functions are applied to arguments of mismatched types, an error will occur.

@jit vs @vectorize​

We've seen the @jit and @vectorize decorators used to optimize various functions. So far, they've been used without an explanation of how they differ. So, what is different about each decorator?

  • @jit is a general-purpose decorator that can optimize a wide variety of different functions.
  • @vectorize is a decorator meant to allow functions that operate on arrays to be written as if they operate on scalars.

This can be visualized with a simple example:

from numba import jit, vectorize, int64
import numpy as np

x = np.arange(5)
y = x + 5

@jit([int64(int64, int64)])
def jitAddArrays(x, y):
    return x + y

@vectorize([int64(int64, int64)])
def vectorizeAddArrays(x, y):
    return x + y

z1 = vectorizeAddArrays(x, y)

print(z1)

z2 = jitAddArrays(x, y)

print(z2)

img

Functions created using @vectorize can be applied to arrays. However, attempting to do the same with @jit results in an error.

Examples​

Numpy matrix​

The following example uses jit to calculate the element-wise sum of a 250x250 matrix. This is a total of 62,500 additions.

Here, the nopython=True option is used. This option produces faster, primitive-optimized code that does not need the Python interpreter to execute. Without this flag, Numba will fall back to the slower object mode in some circumstances.

This example looks at three cases:

  • A regular function without JIT
  • A JIT function that needs compilation
  • A JIT function that is already compiled

The first time the JIT-enabled function is run on a matrix of integers, it's almost ten times slower than the standard function. However, after compilation, the JIT-enabled function is almost two hundred times faster!

from numba import jit
import numpy as np
import time

x = np.arange(62500).reshape(250, 250)
y = 0.01 * x

def calc(a):
    matrix_sum = 0.0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            matrix_sum += a[i, j]
    return matrix_sum

@jit(nopython=True)
def jitCalc(a):
    matrix_sum = 0.0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            matrix_sum += a[i, j]
    return matrix_sum

# Time without JIT
start = time.time()
calc(x)
end = time.time()
print("Execution time (without JIT) = %s" % (end - start))

# Time with compilation and JIT
start = time.time()
jitCalc(x)
end = time.time()
print("Execution time (JIT + compilation) = %s" % (end - start))

# Time with JIT (already compiled)
start = time.time()
jitCalc(x)
end = time.time()
print("Execution time (JIT) = %s" % (end - start))

Using Deephaven tables​

To show how the performance of @jit and @vectorize differ when applied to Deephaven tables, we will create identical functions that use these decorators. We then measure the performance of creating new columns in a 625,000 row table when using the functions.

from deephaven import empty_table
from numba import jit, vectorize, double, int64
import time

def addColumns(A, B):
    return A + B

@jit([int64(int64, int64)])
def jitAddColumns(A, B):
    return A + B

@vectorize([int64(int64, int64)])
def vectorizeAddColumns(A, B):
    return A + B

def cubicFunc(C):
    return 0.0025 * C**3 - 1.75 * C**2 - C + 10

@jit([double(int64)])
def jitCubicFunc(C):
    return 0.0025 * C**3 - 1.75 * C**2 - C + 10

@vectorize([double(int64)])
def vectorizeCubicFunc(C):
    return 0.0025 * C**3 - 1.75 * C**2 - C + 10

t = empty_table(625000).update(formulas=["A = ii", "B = ii"])

# Time column addition without Numba
start = time.time()
t2 = t.update(formulas=["C = addColumns(A, B)"])
end = time.time()
print("column addition - Execution time (without Numba) = %s" % (end - start))

# Time column addition with jit
start = time.time()
t3 = t.update(formulas=["C = jitAddColumns(A, B)"])
end = time.time()
print("column addition - Execution time (jit) = %s" % (end - start))

# Time column addition with vectorize
start = time.time()
t4 = t.update(formulas=["C = vectorizeAddColumns(A, B)"])
end = time.time()
print("column addition - Execution time (vectorize) = %s" % (end - start))

# Time cubic polynomial without Numba
start = time.time()
t5 = t2.update(formulas=["D = cubicFunc(C)"])
end = time.time()
print("cubic polynomial - Execution time (without Numba) = %s" % (end - start))

# Time cubic polynomial with jit
start = time.time()
t5 = t2.update(formulas=["D = jitCubicFunc(C)"])
end = time.time()
print("cubic polynomial - Execution time (jit) = %s" % (end - start))

# Time a cubic polynomial with vectorize
start = time.time()
t6 = t2.update(formulas=["D = vectorizeCubicFunc(C)"])
end = time.time()
print("cubic polynomial - Execution time (vectorize) = %s" % (end - start))

The use of @jit with functions operating on Deephaven tables results in a very small performance increase over its standard counterparts. This performance increase is small enough to make the additional overhead of compiling the function into optimized machine code not worth it.

The use of @vectorize with functions operating on Deephaven tables results in a large performance increase over its standard counterparts. This performance increase is large enough to warrant the additional overhead associated with compiling the function into optimized machine code.