Vectorized Python UDFs¶

Overview¶

Vectorized Python UDFs let you define Python functions that receive batches of input rows as Pandas DataFrames and return batches of results as Pandas arrays or Series. You call vectorized Python UDFs the same way you call other Python UDFs.

Advantages of using vectorized Python UDFs compared to the default row-by-row processing pattern include:

• The potential for better performance if your Python code operates efficiently on batches of rows.

• Less transformation logic required if you are calling into libraries that operate on Pandas DataFrames or Pandas arrays.

When you use vectorized Python UDFs:

• You do not need to change how you write queries using Python UDFs. All batching is handled by the UDF framework rather than your own code.

• As with non-vectorized UDFs, there is no guarantee of which instances of your handler code will see which batches of input.

Getting started with vectorized Python UDFs¶

To create a vectorized Python UDF, use one of the supported mechanisms for annotating your handler function.

CREATE FUNCTION add_one_to_inputs(x NUMBER(10, 0), y NUMBER(10, 0))
  RETURNS NUMBER(10, 0)
  LANGUAGE PYTHON
  RUNTIME_VERSION = 3.9
  PACKAGES = ('pandas')
  HANDLER = 'add_one_to_inputs'
AS $$
import pandas
from _snowflake import vectorized

@vectorized(input=pandas.DataFrame)
def add_one_to_inputs(df):
 return df[0] + df[1] + 1
$$;

CREATE FUNCTION add_one_to_inputs(x NUMBER(10, 0), y NUMBER(10, 0))
  RETURNS NUMBER(10, 0)
  LANGUAGE PYTHON
  RUNTIME_VERSION = 3.9
  PACKAGES = ('pandas')
  HANDLER = 'add_one_to_inputs'
AS $$
import pandas

def add_one_to_inputs(df):
 return df[0] + df[1] + 1

add_one_to_inputs._sf_vectorized_input = pandas.DataFrame
$$;

Setting a target batch size¶

Calls to the Python handler function must execute within a time limit, which is 180 seconds, and each DataFrame passed as input to the handler function may currently contain up to a few thousand rows. In order to stay within the time limit, you may want to set the target batch size for your handler function, which imposes a maximum number of rows per input DataFrame. Note that setting a larger value does not guarantee that Snowflake will encode batches with the specified number of rows. You can set the target batch size using either the vectorized decorator or an attribute on the function.

CREATE FUNCTION add_one_to_inputs(x NUMBER(10, 0), y NUMBER(10, 0))
  RETURNS NUMBER(10, 0)
  LANGUAGE PYTHON
  RUNTIME_VERSION = 3.9
  PACKAGES = ('pandas')
  HANDLER = 'add_one_to_inputs'
AS $$
import pandas
from _snowflake import vectorized

@vectorized(input=pandas.DataFrame, max_batch_size=100)
def add_one_to_inputs(df):
 return df[0] + df[1] + 1
$$;

CREATE FUNCTION add_one_to_inputs(x NUMBER(10, 0), y NUMBER(10, 0))
  RETURNS NUMBER(10, 0)
  LANGUAGE PYTHON
  RUNTIME_VERSION = 3.9
  PACKAGES = ('pandas')
  HANDLER = 'add_one_to_inputs'
AS $$
import pandas

def add_one_to_inputs(df):
 return df[0] + df[1] + 1

add_one_to_inputs._sf_vectorized_input = pandas.DataFrame
add_one_to_inputs._sf_max_batch_size = 100
$$;

DataFrame encoding¶

Batches of arguments to the UDF are encoded as arrays in the input Pandas DataFrames, and the number of rows in each DataFrame may vary. For more information, see Setting a target batch size. Arguments can be accessed in the DataFrame by their index, i.e. the first argument has an index of 0, the second has an index of 1, and so on. The Pandas array or Series that the UDF handler returns must have the same length as that of the input DataFrame.

To illustrate, suppose that you define a vectorized Python UDF as follows:

CREATE OR REPLACE FUNCTION add_inputs(x INT, y FLOAT)
  RETURNS FLOAT
  LANGUAGE PYTHON
  RUNTIME_VERSION = 3.9
  PACKAGES = ('pandas')
  HANDLER = 'add_inputs'
AS $$
import pandas
from _snowflake import vectorized

@vectorized(input=pandas.DataFrame)
def add_inputs(df):
  return df[0] + df[1]
$$;

This UDF uses df[0] to access the Pandas array for the first argument, and df[1] for the second. df[0] + df[1] results in a Pandas array with the pairwise sums of corresponding elements from the two arrays. After creating the UDF, you might call it with some input rows:

SELECT add_inputs(x, y)
FROM (
  SELECT 1 AS x, 3.14::FLOAT as y UNION ALL
  SELECT 2, 1.59 UNION ALL
  SELECT 3, -0.5
);

+------------------+
| ADD_INPUTS(X, Y) |
|------------------|
|             4.14 |
|             3.59 |
|             2.5  |
+------------------+

Here the add_inputs Python function receives a DataFrame analogous to one created with the following Python code:

>>> import pandas
>>> df = pandas.DataFrame({0: pandas.array([1, 2, 3]), 1: pandas.array([3.14, 1.59, -0.5])})
>>> df
   0     1
0  1  3.14
1  2  1.59
2  3 -0.50

The line return df[0] + df[1] in the handler function results in an array similar to the following Python code:

>>> df[0] + df[1]
0    4.14
1    3.59
2    2.50
dtype: float64