Como criar UDFs 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.

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.

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)

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

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

Aqui está a saída:

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

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

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
;

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

Chame a 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
$$;

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.

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

Chame a UDF:

select py_udf();

Aqui está a saída:

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

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.

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

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

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

Chame a UDF:

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

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        |
+-----------------+