Vectorized Python UDTFs

This topic introduces vectorized Python UDTFs.

Overview

Vectorized Python UDTFs (user-defined table functions), which are UDTFs with a vectorized end_partition, enable seamless partition-by-partition processing by operating on partitions as pandas DataFrames and returning results as pandas DataFrames or lists of pandas arrays or pandas Series. This makes for easy integration with libraries that operate on pandas DataFrames or pandas arrays.

Use when:

  • You need to process your data on a partition-by-partition basis instead of on a row-by-row basis.

  • You need to return multiple rows or columns for each partition.

  • You want to use libraries that operate on pandas DataFrames for data analysis.

Prerequisites

The Snowpark Library for Python version 1.6.1 or later is required.

Getting Started

To create a UDTF with a vectorized end_partition:

  • Optionally, define a handler class with an __init__ method which will be invoked before processing each partition.

  • Do not define a process method.

  • Define an end_partition method that takes in a DataFrame argument and returns or yields a pandas.DataFrame or a tuple of pandas.Series or pandas.arrays where each array is a column. The column types of the result must match the column types in the UDTF definition.

  • Mark the end_partition method as vectorized using the @vectorized decorator or the _sf_vectorized_input function attribute. For more information, refer to Vectorized Python UDFs. The @vectorized decorator can only be used when the Python UDTF is executed within Snowflake, for example, when using a SQL worksheet. When you are executing using the client or a Python worksheet, you must use the function attribute.

Note

The default column names for the input DataFrame to a UDTF with a vectorized end_partition match the signature of the SQL function. The column names will follow the SQL identifier requirements. Namely, if an identifier is unquoted it will be capitalized, and if it’s double quoted it will be preserved as it is.

Here is an example of creating a UDTF with a vectorized end_partition, using the @vectorized decorator.

from _snowflake import vectorized
import pandas

class handler:
  def __init__(self):
    # initialize a state
  @vectorized(input=pandas.DataFrame)
  def end_partition(self, df):
    # process the DataFrame
    return result_df
Copy

Here is an example of creating a UDTF with a vectorized end_partition, using the function attribute.

import pandas

class handler:
  def __init__(self):
    # initialize a state
  def end_partition(self, df):
    # process the DataFrame
    return result_df

handler.end_partition._sf_vectorized_input = pandas.DataFrame
Copy

Note

A UDTF with a vectorized end_partition must be called with PARTITION BY clause to build the partitions.

To call the UDTF with all the data in the same partition:

SELECT * FROM table(udtf(x,y,z) OVER (PARTITION BY 1));
Copy

To call the UDTF with the data partitionioned by column x:

SELECT * FROM table(udtf(x,y,z) OVER (PARTITION BY x));
Copy

Type Support

UDTFs with a vectorized end_partition support the same SQL types as Vectorized Python UDFs for arguments and return values. However, for UDTFs with a vectorized end_partition, SQL NUMBER arguments with a scale of 0 that all fit in a 64-bit or smaller integer type will always be mapped to Int16, Int32, or Int64. In other words, unlike scalar UDFs, if the argument of a UDTF is not nullable, it will not be converted to int16, int32, or int64.

To view a table showing how SQL types are mapped to Pandas dtypes, see the type support table in the Vectorized Python UDFs topic.

Example: Row collection using a regular UDTF vs. using a UDTF with a vectorized end_partition

Here is an example of how to do row collection using a regular UDTF.

import pandas

class handler:
  def __init__(self):
    self.rows = []
  def process(self, *row):
    self.rows.append(row)
  def end_partition(self):
    df = pandas.DataFrame(self.rows)
    # process the DataFrame
    return result_df
Copy

Here is an example of how to do row collection using a UDTF with a vectorized end_partition.

from _snowflake import vectorized
import pandas

class handler:
  def __init__(self):
    self.rows = []
  @vectorized(input=pandas.DataFrame)
  def end_partition(self, df):
  # process the DataFrame
    return result_df
Copy

Example: Calculate the summary statistic for each column in the partition

Here is an example of how to calculate the summary statistic for each column in the partition using the pandas describe() method.

First, create a table and generate 3 partitions of 5 rows each.

create or replace table test_values(id varchar, col1 float, col2 float, col3 float, col4 float, col5 float);

-- generate 3 partitions of 5 rows each
insert into test_values
select 'x',
uniform(1.5,1000.5,random(1))::float col1,
uniform(1.5,1000.5,random(2))::float col2,
uniform(1.5,1000.5,random(3))::float col3,
uniform(1.5,1000.5,random(4))::float col4,
uniform(1.5,1000.5,random(5))::float col5
from table(generator(rowcount => 5));

insert into test_values
select 'y',
uniform(1.5,1000.5,random(10))::float col1,
uniform(1.5,1000.5,random(20))::float col2,
uniform(1.5,1000.5,random(30))::float col3,
uniform(1.5,1000.5,random(40))::float col4,
uniform(1.5,1000.5,random(50))::float col5
from table(generator(rowcount => 5));

insert into test_values
select 'z',
uniform(1.5,1000.5,random(100))::float col1,
uniform(1.5,1000.5,random(200))::float col2,
uniform(1.5,1000.5,random(300))::float col3,
uniform(1.5,1000.5,random(400))::float col4,
uniform(1.5,1000.5,random(500))::float col5
from table(generator(rowcount => 5));
Copy

Take a look at the data.

select * from test_values;

-----------------------------------------------------
|"ID"  |"COL1"  |"COL2"  |"COL3"  |"COL4"  |"COL5"  |
-----------------------------------------------------
|x     |8.0     |99.4    |714.6   |168.7   |397.2   |
|x     |106.4   |237.1   |971.7   |828.4   |988.2   |
|x     |741.3   |207.9   |32.6    |640.6   |63.2    |
|x     |541.3   |828.6   |844.9   |77.3    |403.1   |
|x     |4.3     |723.3   |924.3   |282.5   |158.1   |
|y     |976.1   |562.4   |968.7   |934.3   |977.3   |
|y     |390.0   |244.3   |952.6   |101.7   |24.9    |
|y     |599.7   |191.8   |90.2    |788.2   |761.2   |
|y     |589.5   |201.0   |863.4   |415.1   |696.1   |
|y     |46.7    |659.7   |571.1   |938.0   |513.7   |
|z     |313.9   |188.5   |964.6   |435.4   |519.6   |
|z     |328.3   |643.1   |766.4   |148.1   |596.4   |
|z     |929.0   |255.4   |915.9   |857.2   |425.5   |
|z     |612.8   |816.4   |220.2   |879.5   |331.4   |
|z     |487.1   |704.5   |471.5   |378.9   |481.2   |
-----------------------------------------------------
Copy

Next, create the function.

create or replace function summary_stats(id varchar, col1 float, col2 float, col3 float, col4 float, col5 float)
returns table (column_name varchar, count int, mean float, std float, min float, q1 float, median float, q3 float, max float)
language python
runtime_version=3.8
packages=('pandas')
handler='handler'
as $$
from _snowflake import vectorized
import pandas

class handler:
    @vectorized(input=pandas.DataFrame)
    def end_partition(self, df):
      # using describe function to get the summary statistics
      result = df.describe().transpose()
      # add a column at the beginning for column ids
      result.insert(loc=0, column='column_name', value=['col1', 'col2', 'col3', 'col4', 'col5'])
      return result
$$;
Copy

Call the function and partition by id.

-- partition by id
select * from test_values, table(summary_stats(id, col1, col2, col3, col4, col5)
over (partition by id))
order by id, column_name;

--------------------------------------------------------------------------------------------------------------------------------------------------------------------
|"ID"  |"COL1"  |"COL2"  |"COL3"  |"COL4"  |"COL5"  |"COLUMN_NAME"  |"COUNT"  |"MEAN"              |"STD"               |"MIN"  |"Q1"   |"MEDIAN"  |"Q3"   |"MAX"  |
--------------------------------------------------------------------------------------------------------------------------------------------------------------------
|x     |NULL    |NULL    |NULL    |NULL    |NULL    |col1           |5        |280.25999999999993  |339.5609267863427   |4.3    |8.0    |106.4     |541.3  |741.3  |
|x     |NULL    |NULL    |NULL    |NULL    |NULL    |col2           |5        |419.25999999999993  |331.72476995244114  |99.4   |207.9  |237.1     |723.3  |828.6  |
|x     |NULL    |NULL    |NULL    |NULL    |NULL    |col3           |5        |697.62              |384.2964311569911   |32.6   |714.6  |844.9     |924.3  |971.7  |
|x     |NULL    |NULL    |NULL    |NULL    |NULL    |col4           |5        |399.5               |321.2689294033894   |77.3   |168.7  |282.5     |640.6  |828.4  |
|x     |NULL    |NULL    |NULL    |NULL    |NULL    |col5           |5        |401.96000000000004  |359.83584173897964  |63.2   |158.1  |397.2     |403.1  |988.2  |
|y     |NULL    |NULL    |NULL    |NULL    |NULL    |col1           |5        |520.4               |339.16133329139984  |46.7   |390.0  |589.5     |599.7  |976.1  |
|y     |NULL    |NULL    |NULL    |NULL    |NULL    |col2           |5        |371.84              |221.94799616126298  |191.8  |201.0  |244.3     |562.4  |659.7  |
|y     |NULL    |NULL    |NULL    |NULL    |NULL    |col3           |5        |689.2               |371.01012789410476  |90.2   |571.1  |863.4     |952.6  |968.7  |
|y     |NULL    |NULL    |NULL    |NULL    |NULL    |col4           |5        |635.46              |366.6140927460372   |101.7  |415.1  |788.2     |934.3  |938.0  |
|y     |NULL    |NULL    |NULL    |NULL    |NULL    |col5           |5        |594.64              |359.0334218425911   |24.9   |513.7  |696.1     |761.2  |977.3  |
|z     |NULL    |NULL    |NULL    |NULL    |NULL    |col1           |5        |534.22              |252.58182238633088  |313.9  |328.3  |487.1     |612.8  |929.0  |
|z     |NULL    |NULL    |NULL    |NULL    |NULL    |col2           |5        |521.58              |281.4870103574941   |188.5  |255.4  |643.1     |704.5  |816.4  |
|z     |NULL    |NULL    |NULL    |NULL    |NULL    |col3           |5        |667.72              |315.53336907528495  |220.2  |471.5  |766.4     |915.9  |964.6  |
|z     |NULL    |NULL    |NULL    |NULL    |NULL    |col4           |5        |539.8199999999999   |318.73025742781306  |148.1  |378.9  |435.4     |857.2  |879.5  |
|z     |NULL    |NULL    |NULL    |NULL    |NULL    |col5           |5        |470.82              |99.68626786072393   |331.4  |425.5  |481.2     |519.6  |596.4  |
--------------------------------------------------------------------------------------------------------------------------------------------------------------------
Copy

Alternatively, call the function and treat the whole table as one partition.

-- treat the whole table as one partition
select * from test_values, table(summary_stats(id, col1, col2, col3, col4, col5)
over (partition by 1))
order by id, column_name;

---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|"ID"  |"COL1"  |"COL2"  |"COL3"  |"COL4"  |"COL5"  |"COLUMN_NAME"  |"COUNT"  |"MEAN"             |"STD"               |"MIN"  |"Q1"                |"MEDIAN"  |"Q3"    |"MAX"  |
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
|NULL  |NULL    |NULL    |NULL    |NULL    |NULL    |col1           |15       |444.96             |314.01110034974425  |4.3    |210.14999999999998  |487.1     |606.25  |976.1  |
|NULL  |NULL    |NULL    |NULL    |NULL    |NULL    |col2           |15       |437.56             |268.95505944302295  |99.4   |204.45              |255.4     |682.1   |828.6  |
|NULL  |NULL    |NULL    |NULL    |NULL    |NULL    |col3           |15       |684.8466666666667  |331.87254839915937  |32.6   |521.3               |844.9     |938.45  |971.7  |
|NULL  |NULL    |NULL    |NULL    |NULL    |NULL    |col4           |15       |524.9266666666666  |327.074780585783    |77.3   |225.6               |435.4     |842.8   |938.0  |
|NULL  |NULL    |NULL    |NULL    |NULL    |NULL    |col5           |15       |489.14             |288.9176669671038   |24.9   |364.29999999999995  |481.2     |646.25  |988.2  |
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Copy

Best Practices

This section describes best practices.

  1. To improve performance and prevent timeouts, avoid using pandas.concat to accumulate partial results. Instead, yield the partial result whenever one is ready. For example, instead of:

    results = []
    while(...):
      partial_result = pd.DataFrame(...)
      results.append(partial_result)
    return pd.concat(results)
    
    Copy

    Do this:

    while(...):
      partial_result = pd.DataFrame(...)
      yield partial_result
    
    Copy
  2. If a scalar must be returned with each row, build a list of repeated values instead of unpackaging the numpy array to create tuples. For example, for a 2-column result, instead of:

    return tuple(map(lambda n: (scalar_value, n[0], n[1]), results))
    
    Copy

    Do this:

    return tuple([scalar_value] * len(results), results[:, 0], results[:, 1])
    
    Copy
  3. To improve performance, unpackage semi-structured data into columns. For example, if you have a variant column, obj, with elements, x(int), y(float), and z(string), then instead of defining a UDTF with a signature like this:

    create function vec_udtf(variant obj)
    
    Copy

    And calling it using vec_udtf(obj), you should define the UDTF with signature:

    create function vec_udtf(int, float, string)
    
    Copy

    And call it using vec_udtf(obj:x, obj:y, obj:z).

  4. By default, Snowflake encodes the inputs into pandas dtypes that support NULL values (for example, Int64). If you are using a library that requires a primitive type (such as numpy) and your input has no NULL values, you should cast the column to a primitive type before using the library. For example:

    input_df['y'] =  input_df['y'].astype("int64")
    
    Copy

    For more information, see Type Support.