Creating Python UDFs¶

Writing the Python Module and Function¶

Write a module that follows the specifications below:

• Define the module. A module is a file containing Python definitions and statements.

• Define a function inside the module.

• If the function accepts arguments, each argument must be one of the data types specified in the Python Data Type column of the SQL-Python Type Mappings table.

  Function arguments are bound by position, not name. The first argument passed to the UDF is the first argument received by the Python function.

• Specify an appropriate return value. Because a Python UDF must be a scalar function, it must return one value each time that it is invoked. The type of the return value must be one of the data types specified in the Python Data Type column of the SQL-Python Type Mappings table. The type of the return value must be compatible with the SQL data type specified in the RETURNS clause of the CREATE FUNCTION statement.

• Your module can contain more than one function. The function that is called by Snowflake can call other functions in the same module, or in other modules.

• Your function (and any functions called by your function) must comply with the Snowflake-imposed constraints for Python UDFs.

Reading and Writing Files with a UDF Handler¶

You can read and write files with UDF handler code. To do this securely within the restricted engine in which Snowflake executes UDFs, follow the guidelines described here.

def func(text):
   # Ensure the file name's uniqueness by appending the function's process ID.
   file_path = '/tmp/content' + str(os.getpid())
   with open(file_path, "w") as file:
      file.write(text)

Creating an In-line Python UDF¶

For an in-line UDF, you supply the Python source code as part of the CREATE FUNCTION statement.

For example, the following statement creates an in-line Python UDF that adds one to a given integer:

create or replace function addone(i int)
returns int
language python
runtime_version = '3.8'
handler = 'addone_py'
as
$$
def addone_py(i):
  return i+1
$$;

The Python source code is specified in the AS clause. The source code can be surrounded by either single quotes or by a pair of dollar signs ($$). Using the double dollar signs is usually easier if the source code contains embedded single quotes.

Call the UDF:

select addone(10);

Here is the output:

+------------+
| ADDONE(10) |
|------------|
|         11 |
+------------+

The Python source code can contain more than one module, and more than one function in a module, so the HANDLER clause specifies the module and function to call.

An in-line Python UDF can call code in modules that are included in the IMPORTS clause.

For more details about the syntax of the CREATE FUNCTION statement, see CREATE FUNCTION.

For more examples, see in-line Python UDF examples.

Creating a Python UDF With Code Uploaded from a Stage¶

The following statements create a simple Python UDF using code uploaded from a stage. The stage hosting the file must be readable by the owner of the UDF. Also, ZIP files must be self-contained and not rely on any additional setup scripts to be executed.

Create a Python file named sleepy.py that contains your source code:

def snore(n):   # return a series of n snores
    result = []
    for a in range(n):
        result.append("Zzz")
    return result

Launch the SnowSQL (CLI Client) and use the PUT command to copy the file from the local file system to the default user stage, named @~. (The PUT command cannot be executed through the Snowflake GUI.)

put
file:///Users/Me/sleepy.py
@~/
auto_compress = false
overwrite = true
;

If you delete or rename the file, you can no longer call the UDF. If you need to update your file, then update it while no calls to the UDF can be made. If the old file is still in the stage, the PUT command should include the clause OVERWRITE=TRUE.

Create the UDF. The handler specifies the module and the function.

create or replace function dream(i int)
returns variant
language python
runtime_version = '3.8'
handler = 'sleepy.snore'
imports = ('@~/sleepy.py')

Call the UDF:

select dream(3);

+----------+
| DREAM(3) |
|----------|
| [        |
|   "Zzz", |
|   "Zzz", |
|   "Zzz"  |
| ]        |
+----------+

create or replace function multiple_import_files(s string)
returns string
language python
runtime_version=3.8
imports=('@python_udf_dep/bar/python_imports_a.zip', '@python_udf_dep/foo/python_imports_b.zip')
handler='compute'
as
$$
def compute(s):
  return s
$$;

Using an Imported Package in an In-line Python UDF¶

A curated list of 3rd party packages from Anaconda is available. For more information, see Using Third-Party Packages.

The following code shows how to import packages and return their versions.

Create the UDF:

create or replace function py_udf()
returns variant
language python
runtime_version = 3.8
packages = ('numpy','pandas','xgboost==1.5.0')
handler = 'udf'
as $$
import numpy as np
import pandas as pd
import xgboost as xgb
def udf():
    return [np.__version__, pd.__version__, xgb.__version__]
$$;

Call the UDF:

select py_udf();

Here is the output:

+-------------+
| PY_UDF()    |
|-------------|
| [           |
|   "1.19.2", |
|   "1.4.0",  |
|   "1.5.0"   |
| ]           |
+-------------+

Loading a File from a Stage into a Python UDF¶

This example shows how to import non-code files into a Python UDF from a stage. The file will only be read once during UDF creation, and will not be read again during UDF execution if reading the file happens outside of the target handler. For more about reading a file, see Reading Files with a UDF Handler.

Create the UDF:

create or replace function my_udf()
  returns string
  language python
  runtime_version=3.8
  imports=('@my_stage/file.txt')
  handler='compute'
as
$$
import sys
import os

with open(os.path.join(sys._xoptions["snowflake_import_directory"], 'file.txt'), "r") as f:
    s = f.read()

def compute():
    return s
$$;

Unzipping a Staged File¶

You can store a .zip file on a stage, then unzip it in a UDF by using the Python zipfile module.

For example, you can upload a .zip file to a stage, then reference the .zip file at its staged location in the IMPORTS clause when you create the UDF. At run time, Snowflake will copy the staged file into an import directory from which your code can access it.

For more about reading and writing files, see Reading and Writing Files with a UDF Handler.

In the following example, the UDF code uses an NLP model to discover entities in text. The code returns an array of these entities. To set up the NLP model for processing the text, the code first uses the zipfile module to extract the file for the model (en_core_web_sm-2.3.1) from a .zip file. The code then uses the spaCy module to load the model from the file.

Note that the code writes extracted file contents to the /tmp directory created for the query calling this function. The code uses file locks to ensure that the extraction is synchronized across Python worker processes; this way, contents are unzipped only once. For more about writing files, see Writing Files with a UDF Handler.

For more about the zipfile module, see the zipfile reference. For more about the spaCy module, see the spaCy API documentation.

Create the UDF:

create or replace function py_spacy(str string)
   returns array
   language python
   runtime_version = 3.8
   handler = 'func'
   packages = ('spacy')
   imports = ('@spacy_stage/spacy_en_core_web_sm.zip')
   as
$$
import fcntl
import os
import spacy
import sys
import threading
import zipfile

# File lock class for synchronizing write access to /tmp
class FileLock:
   def __enter__(self):
      self._lock = threading.Lock()
      self._lock.acquire()
      self._fd = open('/tmp/lockfile.LOCK', 'w+')
      fcntl.lockf(self._fd, fcntl.LOCK_EX)

   def __exit__(self, type, value, traceback):
      self._fd.close()
      self._lock.release()

# Get the location of the import directory. Snowflake sets the import
# directory location so code can retrieve the location via sys._xoptions.
IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]

# Get the path to the ZIP file and set the location to extract to.
zip_file_path = import_dir + "spacy_en_core_web_sm.zip"
extracted = '/tmp/en_core_web_sm'

# Extract the contents of the ZIP. This is done under the file lock
# to ensure that only one worker process unzips the contents.
with FileLock():
   if not os.path.isdir(extracted + '/en_core_web_sm/en_core_web_sm-2.3.1'):
      with zipfile.ZipFile(zip_file_path, 'r') as myzip:
         myzip.extractall(extracted)

# Load the model from the extracted file.
nlp = spacy.load(extracted + "/en_core_web_sm/en_core_web_sm-2.3.1")

def func(text):
   doc = nlp(text)
   result = []

   for ent in doc.ents:
      result.append((ent.text, ent.start_char, ent.end_char, ent.label_))
   return result
$$;

NULL Handling in Python UDFs¶

The following code shows how NULL values are handled. For more information, see NULL Values.

Create the UDF:

create or replace function py_udf_null(a variant)
returns string
language python
runtime_version = 3.8
handler = 'udf'
as $$

def udf(a):
    if not a:
        return 'JSON null'
    elif getattr(a, "is_sql_null", False):
        return 'SQL null'
    else:
        return 'not null'
$$;

Call the UDF:

select py_udf_null(null);
select py_udf_null(parse_json('null'));
select py_udf_null(10);

Here is the output:

+-------------------+
| PY_UDF_NULL(NULL) |
|-------------------|
| SQL null          |
+-------------------+

+---------------------------------+
| PY_UDF_NULL(PARSE_JSON('NULL')) |
|---------------------------------|
| JSON null                       |
+---------------------------------+

+-----------------+
| PY_UDF_NULL(10) |
|-----------------|
| not null        |
+-----------------+