Como criar UDFs de Python

Este tópico mostra como criar e instalar uma UDF (função definida pelo usuário) de Python.

Neste tópico:

Como escrever o código de Python

Como escrever o módulo e a função de Python

Escreva um módulo que siga as especificações abaixo:

  • Defina o módulo. Um módulo é um arquivo contendo definições e instruções de Python.

  • Defina uma função dentro do módulo.

  • Se a função aceitar argumentos, cada argumento deve ser um dos tipos de dados especificados na coluna Python Data Type da tabela de mapeamento entre tipos de SQL e Python.

    Os argumentos de função estão vinculados à posição, não ao nome. O primeiro argumento passado para a UDF é o primeiro argumento recebido pela função de Python.

  • Especifique um valor de retorno adequado. Como uma UDF de Python deve ser uma função escalar, ela deve retornar um valor cada vez que for invocada. O tipo do valor de retorno deve ser um dos tipos de dados especificados na coluna Python Data Type da tabela de mapeamento entre tipos de SQL e Python. O tipo do valor de retorno deve ser compatível com o tipo de dados de SQL especificado na cláusula RETURNS da instrução CREATE FUNCTION.

  • Seu módulo pode conter mais de uma função. A função que é chamada pelo Snowflake pode chamar outras funções no mesmo módulo ou em outros módulos.

  • Sua função (e quaisquer funções chamadas por sua função) devem estar de acordo com as restrições impostas pelo Snowflake para UDFs de Python.

Nota

Há uma API de lote de UDF de Python que permite definir funções de Python que recebem lotes de linhas de entrada como DataFrames do Pandas e retornam lotes de resultados como Arrays ou Series do Pandas. Para obter mais informações, consulte API de lote de UDF de Python.

Como escrever e ler arquivos com um manipulador de UDF

Você pode ler e escrever arquivos com o código do manipulador da UDF. Para fazer isso com segurança dentro do mecanismo restrito no qual o Snowflake executa UDFs, siga as orientações descritas aqui.

Como ler arquivos com um manipulador de UDF

Um manipulador de UDF pode ler arquivos que foram carregados para um estágio do Snowflake. O estágio que hospeda o arquivo deve ser legível pelo proprietário da UDF.

Quando você especifica a localização do estágio de um arquivo na cláusula IMPORTS do CREATE FUNCTION, o Snowflake copia o arquivo preparado para um diretório de importação disponível especificamente para a UDF. Seu código do manipulador pode ler o arquivo a partir daí.

Note que o Snowflake copia todos os arquivos importados, potencialmente de vários estágios, para um único diretório de importação. Por essa razão, os nomes dos arquivos especificados na cláusula IMPORTS devem ser distintos uns dos outros.

Para obter um código de exemplo, consulte Como carregar um arquivo de um estágio em uma UDF de Python neste tópico.

Para ler um arquivo com o código do manipulador da UDF:

  1. Copie o arquivo para o estágio do Snowflake.

    Você pode usar o comando PUT para carregar arquivos de um diretório local em um computador cliente; para obter mais informações, consulte PUT. Para obter informações mais gerais sobre o carregamento de arquivos para um estágio, consulte Visão geral do carregamento de dados.

  2. Ao criar a UDF usando CREATE FUNCTION, especifique a localização do arquivo na cláusula IMPORTS.

    O código no exemplo a seguir especifica um arquivo file.txt em um estágio chamado my_stage.

    create or replace function my_udf()
       ...
       imports=('@my_stage/file.txt')
       ...
    
    Copy
  3. Em seu código do manipulador, leia o arquivo do diretório de importação.

    O Snowflake copia o arquivo preparado para o diretório de importação da UDF. Você pode recuperar a localização do diretório usando a opção snowflake_import_directory do sistema.

    Em código de Python, você pode recuperar a localização do diretório usando o método de Python sys._xoptions, como no exemplo a seguir:

    IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
    import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
    
    def compute():
       with open(import_dir + 'file.txt', 'r') as file:
          return file.read()
    
    Copy

Como escrever arquivos com um manipulador de UDF

Um manipulador de UDF pode escrever arquivos em um diretório /tmp criado para a consulta chamando a UDF.

Tenha em mente que um diretório /tmp é reservado para uma única consulta de chamada, mas vários processos de trabalhadores de Python podem estar em execução ao mesmo tempo. Para evitar colisões, você deve garantir que o acesso ao diretório /tmp seja sincronizado com outros processos de trabalhadores de Python ou que os nomes dos arquivos escritos no /tmp sejam únicos.

Para obter um código de exemplo, consulte Como descompactar um arquivo preparado neste tópico.

No exemplo a seguir, o código escreve a entrada text no diretório /tmp, anexando o processo da função ID para garantir a exclusividade do local do arquivo.

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)
Copy

Criação da função no Snowflake

Você deve executar uma instrução CREATE FUNCTION para especificar:

  • O nome da função de SQL a ser usada.

  • O nome da função de Python a ser chamada quando a UDF de Python é chamada.

O nome da UDF não precisa corresponder ao nome da função do manipulador escrita em Python. A instrução CREATE FUNCTION associa o nome da UDF com a função de Python.

Ao escolher um nome para a UDF:

  • Siga as regras para Identificadores de objetos.

  • Escolha um nome único ou siga as regras para Sobrecarga.

    Importante

    Ao contrário da sobrecarga para UDFs de SQL, que distingue entre funções baseadas tanto no número quanto nos tipos de dados dos argumentos, UDFs de Python distinguem entre funções com base somente no número de argumentos.

Os argumentos de função estão vinculados à posição, não ao nome. O primeiro argumento passado para a UDF é o primeiro argumento recebido pela função de Python.

Para obter mais informações sobre os tipos de dados de argumentos, consulte Mapeamentos de tipos de dados SQL-Python.

UDFs com código inline vs. UDFs com código carregado de um estágio

O código para uma UDF de Python pode ser especificado de uma das seguintes maneiras:

  • Carregado de um estágio: a instrução CREATE FUNCTION especifica a localização de um código-fonte de Python existente em um estágio.

  • Inline: a instrução CREATE FUNCTION especifica o código-fonte de Python.

Como criar uma UDF de Python inline

Para criar uma UDF inline, forneça o código-fonte de Python como parte da instrução CREATE FUNCTION.

Por exemplo, a seguinte instrução cria uma UDF de Python inline que adiciona um a um determinado número inteiro:

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
$$;
Copy

O código-fonte de Python é especificado na cláusula AS. O código-fonte pode ser colocado entre aspas simples ou entre um par de cifrões ($$). O uso dos cifrões geralmente é mais fácil se o código-fonte contiver aspas simples incorporadas.

Chame a UDF:

select addone(10);
Copy

Aqui está a saída:

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

O código-fonte de Python pode conter mais de um módulo, e mais de uma função em um módulo. Portanto, a cláusula HANDLER especifica o módulo e a função a ser chamada.

Uma UDF de Python inline pode chamar código em módulos que estão incluídos na cláusula IMPORTS.

Para obter mais detalhes sobre a sintaxe da instrução CREATE FUNCTION, consulte CREATE FUNCTION.

Para obter mais exemplos, consulte exemplos de UDF de Python inline.

Como criar uma UDF de Python com código carregado de um estágio

As instruções a seguir criam uma UDF de Python simples usando código carregado de um estágio. O estágio que hospeda o arquivo deve ser legível pelo proprietário da UDF. Além disso, arquivos ZIP devem ser autocontidos e não confiar em nenhum script de configuração adicional para serem executados.

Crie um arquivo de Python chamado sleepy.py contendo seu código-fonte:

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

Inicie o comando SnowSQL (cliente CLI) e use o comando PUT para copiar o arquivo do sistema de arquivos local para o estágio padrão do usuário, chamado @~. (O comando PUT não pode ser executado pela GUI do Snowflake).

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

Se você apagar ou renomear o arquivo, você não poderá mais chamar a UDF. Se você precisar atualizar seu arquivo, atualize-o enquanto nenhuma chamada para a UDF possa ser feita. Se o arquivo antigo ainda estiver no estágio, o comando PUT deverá incluir a cláusula OVERWRITE=TRUE.

Crie a UDF. O manipulador especifica o módulo e a função.

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

Chame a UDF:

select dream(3);

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

Como especificar múltiplos arquivos de importação

Aqui está um exemplo de como especificar múltiplos arquivos de importação.

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
$$;
Copy

Nota

Os nomes dos arquivos de importação especificados devem ser diferentes. Por exemplo, isso não vai funcionar: imports=('@python_udf_dep/bar/python_imports.zip', '@python_udf_dep/foo/python_imports.zip').

Concessão de privilégios sobre a função

Para que alguém com qualquer função que não seja a de proprietário chame a função, o proprietário deve conceder os privilégios adequados para a função.

As instruções GRANT para uma UDF de Python são essencialmente idênticas às instruções GRANT para outras UDFs, tais como UDFs de JavaScript.

Por exemplo:

GRANT USAGE ON FUNCTION my_python_udf(number, number) TO my_role;
Copy

Exemplos

Como usar um pacote importado em uma UDF de Python inline

Uma lista curada de pacotes de terceiros do Anaconda está disponível. Para obter mais informações, consulte Como usar pacotes de terceiros.

Nota

Antes que você possa usar os pacotes fornecidos pelo Anaconda, seu administrador da organização Snowflake deve reconhecer os Termos de Terceiros do Snowflake. Para obter mais informações, consulte Introdução.

O código a seguir mostra como importar pacotes e retornar suas versões.

Crie a 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__]
$$;
Copy

Chame a UDF:

select py_udf();
Copy

Aqui está a saída:

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

Como carregar um arquivo de um estágio em uma UDF de Python

Este exemplo mostra como importar arquivos não codificados para uma UDF de Python a partir de um estágio. O arquivo só será lido uma vez durante a criação do UDF, e não será lido novamente durante a execução do UDF se a leitura do arquivo for feita fora do manipulador de destino. Para saber mais sobre como ler um arquivo, consulte Como ler arquivos com um manipulador de UDF.

Nota

Você só pode importar arquivos do diretório de nível superior em um estágio, não em subpastas.

Crie a 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
$$;
Copy

Como descompactar um arquivo preparado

Você pode armazenar um arquivo .zip em um estágio e depois descompactá-lo em uma UDF usando o módulo Python zipfile.

Por exemplo, você pode carregar um arquivo .zip para um estágio, depois referenciar o arquivo .zip em seu local de estágio na cláusula IMPORTS ao criar a UDF. No runtime, o Snowflake copiará o arquivo preparado para um diretório de importação do qual seu código pode acessá-lo.

Para saber mais sobre como ler e escrever arquivos, consulte Como escrever e ler arquivos com um manipulador de UDF.

No exemplo a seguir, o código da UDF usa um modelo de NLP para descobrir entidades no texto. O código retorna um conjunto dessas entidades. Para configurar o modelo de NLP para processar o texto, o código primeiro usa o módulo zipfile para extrair o arquivo para o modelo (en_core_web_sm-2.3.1) de um arquivo .zip. O código então usa o módulo spaCy para carregar o modelo a partir do arquivo.

Note que o código escreve o conteúdo do arquivo extraído no diretório /tmp criado para a consulta que chama essa função. O código usa bloqueios de arquivo para garantir que a extração seja sincronizada entre os processos de trabalhadores de Python; dessa forma, o conteúdo é descompactado apenas uma vez. Para obter mais informações sobre como escrever arquivos, consulte Como escrever arquivos com um manipulador de UDF.

Para saber mais sobre o módulo de arquivo zip, consulte a referência zipfile. Para saber mais sobre o módulo spaCy, consulte a documentação spaCy API.

Crie a 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
$$;
Copy

Tratamento de valores NULL em UDFs de Python

O código a seguir mostra como os valores NULL são tratados. Para obter mais informações, consulte Valores NULL.

Crie a 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'
$$;
Copy

Chame a UDF:

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

Aqui está a saída:

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

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

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