Writing a UDTF in Python¶

Defining a process method¶

Define a process method that receives as values the UDTF arguments converted from SQL types, returning data that Snowflake will use to create the UDTF’s tabular return value.

The method’s signature must be of the following form:

def process(self, *args):

Your process method must:

• Have a self parameter.

• Declare method parameters corresponding to UDTF parameters.

  Method parameter names needn’t match UDTF parameter names, but the method parameters must be declared in the same order as UDTF parameters are declared.

  When passing UDTF argument values to your method, Snowflake will convert the values from SQL types to the Python types you use in the method. For information about how Snowflake maps between SQL and Python data types, see SQL-Python Data Type Mappings.

• Yield one or more tuples (or return an iterable containing tuples), in which the sequence of tuples corresponds to the sequence of UDTF return value columns.

  The tuple elements must appear in the same order as UDTF return value columns are declared. For more information, see Returning a value in this topic.

  Snowflake will convert values from Python types to SQL types required by the UDTF declaration. For information about how Snowflake maps between SQL and Python data types, see SQL-Python Data Type Mappings.

If a method in the handler class throws an exception, processing will stop. The query that called the UDTF will fail with an error message. If the process method returns None, processing stops. (The end_partition method is still invoked even if the process method returns None.)

process Method Example

Code in the following example shows a StockSale handler class with a process method that processes three UDTF arguments (symbol, quantity, and price), returning a single row with two columns (symbol and total). Note that process method parameters are declared in the same order as stock_sale parameters. Arguments in the process method’s yield statement are in the same order as columns declared in the stock_sale RETURNS TABLE clause.

create or replace function stock_sale(symbol varchar, quantity number, price number(10,2))
returns table (symbol varchar, total number(10,2))
language python
RUNTIME_VERSION = 3.9
handler='StockSale'
as $$
class StockSale:
    def process(self, symbol, quantity, price):
      cost = quantity * price
      yield (symbol, cost)
$$;

Code in the following example calls the preceding UDF, passing values from columns symbol, quantity, and price from the stocks_table table.

select stock_sale.symbol, total
  from stocks_table, table(stock_sale(symbol, quantity, price) over (partition by symbol));

Returning a value¶

When returning output rows, you can use either yield or return (but not both) to return tuples with the tabular value. If the method returns or yields None, processing for the current row stops.

• When using yield, execute a separate yield statement for each output row. This is the best practice because the lazy evaluation that comes with yield enables more efficient processing and can help avoid timeouts.

  Each element in the tuple becomes a column value in the result returned by the UDTF. The order of yield arguments must match the order of columns declared for the return value in the RETURNS TABLE clause of CREATE FUNCTION.

  Code in the following example returns values representing two rows.

  def process(self, symbol, quantity, price):
    cost = quantity * price
    yield (symbol, cost)
    yield (symbol, cost)
  

  Copy

  Note that because the yield argument is a tuple, you must include a trailing comma when passing a single value in the tuple, as in the following example.

  yield (cost,)
  

  Copy

• When using return, return an iterable with tuples.

  Each value in a tuple becomes a column value in the result returned by the UDTF. The order of column values in a tuple must match the order of columns declared for the return value in the RETURNS TABLE clause of CREATE FUNCTION.

  Code in the following example returns two rows, each with two columns: symbol and total.

  def process(self, symbol, quantity, price):
    cost = quantity * price
    return [(symbol, cost), (symbol, cost)]
  

  Copy

Skipping rows¶

To skip an input row and process the next row (such as when you’re validating the input rows), have the process method return one of the following:

• When using return, return None, a list containing None, or an empty list to skip the row.

• When using yield, return None to skip a row.

  Note that if you have multiple calls to yield, any calls after a call that returns None will be ignored by Snowflake.

Code in the following example returns only the rows for which number is a positive integer. If number is not positive, the method returns None to skip the current row and continue processing the next row.

def process(self, number):
  if number < 1:
    yield None
  else:
    yield (number)

Stateful and stateless processing¶

You can implement the handler to process rows in a partition-aware manner or to process them simply row by row.

• In partition-aware processing, the handler includes code to manage partition-scoped state. This includes an __init__ method that executes at the start of partition processing and an end_partition method that Snowflake invokes after processing the partition’s last row. For more information, see Processing partitions in this topic.

• In partition-unaware processing, the handler executes statelessly, ignoring partition boundaries.

  To have the handler execute this way, do not include an __init__ or end_partition method.

Finalizing partition processing¶

You can optionally implement an end_partition method in your handler class that Snowflake will invoke after you have processed all rows in a partition. In this method, you can execute code for a partition after all of the partition’s rows have been processed. Your end_partition method may produce output rows, such as to return the results of a partition-scoped calculation. For more information, see Processing partitions in this topic.

The method’s signature must be of the following form:

def end_partition(self):

Snowflake expects the following of your end_partition method implementation:

• It must not be static.

• It may not have any parameters other than self.

• As an alternative to returning a tabular value, it may produce an empty list or None.

Partition handling example¶

Code in the following example calculates the total cost paid across purchases for a stock by first calculating the cost per purchase and adding purchases together (in the process method). The code returns the total in the end_partition method.

For an example of a UDTF that includes this handler, along with calling the UDTF, refer to Handler class example.

class StockSaleSum:
  def __init__(self):
    self._cost_total = 0
    self._symbol = ""

  def process(self, symbol, quantity, price):
    self._symbol = symbol
    cost = quantity * price
    self._cost_total += cost
    yield (symbol, cost)

  def end_partition(self):
    yield (self._symbol, self._cost_total)

When processing partitions, keep in mind the following:

• Your code may handle partitions that aren’t explicitly specified in a call to the UDTF. Even when a call to the UDTF doesn’t include a PARTITION BY clause, Snowflake partitions the data implicitly.

• Your process method will receive row data in the order specified by the partition’s ORDER BY clause, if any.