Création de fonctions définies par l’utilisateur (UDFs) pour DataFrames dans Python¶

Introduction¶

Avec Snowpark, vous pouvez créer des fonctions définies par l’utilisateur (UDFs) pour vos Lambdas et fonctions personnalisées, et vous pouvez appeler ces UDFs pour traiter les données dans votre DataFrame.

Lorsque vous utilisez l’API de Snowpark pour créer une UDF, la bibliothèque Snowpark charge le code de votre fonction dans une zone de préparation interne. Lorsque vous appelez l’UDF, la bibliothèque Snowpark exécute votre fonction sur le serveur, où se trouvent les données. Par conséquent, les données ne doivent pas être transférées au client pour que la fonction puisse les traiter.

Dans votre code personnalisé, vous pouvez également importer des modules à partir de fichiers Python ou de paquets tiers.

Vous pouvez créer une UDF pour votre code personnalisé de deux façons :

• Vous pouvez créer une UDF anonyme et affecter la fonction à une variable. Tant que cette variable est dans le scope, vous pouvez l’utiliser pour appeler l’UDF.

• Vous pouvez créer une UDF nommée et appeler l’UDF par son nom. Vous pouvez l’utiliser si, par exemple, vous devez appeler une UDF par son nom ou utiliser l’UDF dans une session ultérieure.

Les sections suivantes expliquent comment créer ces UDFs à l’aide d’un environnement de développement local ou à l’aide d’une feuille de calcul Python.

Notez que si vous avez défini une UDF en exécutant la commande CREATE FUNCTION, vous pouvez appeler cette UDF dans Snowpark. Pour plus de détails, voir Appel de fonctions définies par l’utilisateur (UDFs).

import snowflake.snowpark as snowpark
from snowflake.snowpark.functions import col

def main(session: snowpark.Session):
   df_table = session.table("sample_product_data")

Spécifier des dépendances pour une UDF¶

Pour définir une UDF à l’aide de l’API Snowpark, vous devez importer les fichiers contenant tous les modules dont dépend votre UDF, tels que les fichiers Python, les fichiers zip, les fichiers de ressources, etc.

• Pour réaliser cette opération à l’aide de feuilles de calcul Python, reportez-vous à Ajouter un fichier Python d’une zone de préparation à une feuille de calcul.

• Pour réaliser cette opération à l’aide de votre environnement de développement local, vous devez appeler Session.add_import() dans votre code.

Vous pouvez également spécifier un répertoire et la bibliothèque Snowpark le compresse automatiquement et le charge sous forme de fichier zip. (Pour plus de détails sur la lecture des ressources d’une UDF, voir Lecture de fichiers à partir d’une UDF).

Lorsque vous appelez Session.add_import(), la bibliothèque Snowpark charge les fichiers spécifiés dans une zone de préparation interne et importe les fichiers lors de l’exécution de votre UDF.

L’exemple suivant montre comment ajouter un fichier zip dans une zone de préparation en tant que dépendance à votre code :

# Add a zip file that you uploaded to a stage.
session.add_import("@my_stage/<path>/my_library.zip")

Les exemples suivants montrent comment ajouter un fichier Python à partir de votre machine locale :

# Import a Python file from your local machine.
session.add_import("/<path>/my_module.py")

# Import a Python file from your local machine and specify a relative Python import path.
session.add_import("/<path>/my_module.py", import_path="my_dir.my_module")

Les exemples suivants montrent comment ajouter d’autres types de dépendances :

# Add a directory of resource files.
session.add_import("/<path>/my-resource-dir/")

# Add a resource file.
session.add_import("/<path>/my-resource.xml")

Vous n’avez pas besoin de spécifier les dépendances suivantes :

• Vos bibliothèques intégrées Python.

  Ces bibliothèques sont déjà disponibles dans l’environnement d’exécution sur le serveur où vos UDFs sont exécutées.

Utilisation de paquets tiers à partir d’Anaconda dans une UDF¶

Vous pouvez utiliser des packages tiers du canal Snowflake Anaconda dans une UDF.

• Si vous créez une UDF Python dans une feuille de calcul Python, les packages Anaconda sont déjà disponibles pour votre feuille de calcul. Reportez-vous à Ajouter un fichier Python d’une zone de préparation à une feuille de calcul.

• Si vous créez une UDF Python dans votre environnement de développement local, vous pouvez spécifier les packages Anaconda à installer.

Lorsque des requêtes qui appellent des UDFs Python sont exécutées dans un entrepôt Snowflake, les paquets Anaconda sont installés de manière transparente et mis en cache dans l’entrepôt virtuel en votre nom.

Pour plus d’informations sur les meilleures pratiques, sur la façon de visualiser les paquets disponibles et sur la façon de configurer un environnement de développement local, voir Utilisation de paquets tiers.

Si vous écrivez une UDF Python dans votre environnement de développement local, utilisez session.add_packages pour ajouter des packages au niveau de la session.

Cet exemple de code montre comment importer des paquets et renvoyer leurs versions.

import numpy as np
import pandas as pd
import xgboost as xgb
from snowflake.snowpark.functions import udf

session.add_packages("numpy", "pandas", "xgboost==1.5.0")

@udf
def compute() -> list:
   return [np.__version__, pd.__version__, xgb.__version__]

Vous pouvez également utiliser session.add_requirements pour spécifier des paquets avec un fichier d’exigences.

session.add_requirements("mydir/requirements.txt")

Vous pouvez ajouter les paquets de niveau UDF pour remplacer les packages de niveau session que vous avez pu ajouter précédemment.

import numpy as np
import pandas as pd
import xgboost as xgb
from snowflake.snowpark.functions import udf

@udf(packages=["numpy", "pandas", "xgboost==1.5.0"])
   def compute() -> list:
   return [np.__version__, pd.__version__, xgb.__version__]

Création d’une UDF anonyme¶

Pour créer une UDF anonyme, vous pouvez soit :

• Appeler la fonction udf dans le module snowflake.snowpark.functions et transmettre la définition de la fonction anonyme.

• Appeler la méthode register dans la classe UDFRegistration, et transmettre la définition de la fonction anonyme.

Voici un exemple d’une UDF anonyme :

from snowflake.snowpark.types import IntegerType
from snowflake.snowpark.functions import udf

add_one = udf(lambda x: x+1, return_type=IntegerType(), input_types=[IntegerType()])

Création et enregistrement d’une UDF nommée¶

Si vous souhaitez appeler une UDF par son nom (par exemple, en utilisant la fonction call_udf dans le module functions), vous pouvez créer et enregistrer une UDF nommée. Pour ce faire, voici ce que vous pouvez utiliser :

• La méthode register dans la classe UDFRegistration avec l’argument name.

• La fonction udf, dans le module snowflake.snowpark.functions avec l’argument name.

Pour accéder à un attribut ou à une méthode de la classe UDFRegistration appelez la propriété udf de la classe Session.

L’appel d’un register ou d’une udf créera une UDF temporaire que vous pouvez utiliser dans la session en cours.

Pour créer une UDF permanente, appelez la méthode register ou la fonction udf et définissez l’argument is_permanent sur True. Lorsque vous créez une UDF permanente, vous devez également définir l’argument stage_location à l’emplacement de la zone de préparation où le fichier Python pour l’UDF et ses dépendances sont chargées.

Voici un exemple de la manière d’enregistrer une UDF temporaire nommée :

from snowflake.snowpark.types import IntegerType
from snowflake.snowpark.functions import udf

add_one = udf(lambda x: x+1, return_type=IntegerType(), input_types=[IntegerType()], name="my_udf", replace=True)

Voici un exemple de la façon d’enregistrer une UDF permanente nommée en définissant l’argument is_permanent sur True :

@udf(name="minus_one", is_permanent=True, stage_location="@my_stage", replace=True)
def minus_one(x: int) -> int:
   return x-1

Voici un exemple d’appel de ces UDFs :

df = session.create_dataframe([[1, 2], [3, 4]]).to_df("a", "b")
df.select(add_one("a"), minus_one("b")).collect()

[Row(MY_UDF("A")=2, MINUS_ONE("B")=1), Row(MY_UDF("A")=4, MINUS_ONE("B")=3)]

Vous pouvez également appeler l’UDF en utilisant SQL :

session.sql("select minus_one(1)").collect()

[Row(MINUS_ONE(1)=0)]

Création d’une UDF à partir d’un fichier source Python¶

Si vous créez votre UDF dans votre environnement de développement local, vous pouvez également définir votre gestionnaire d’UDF dans un fichier Python, puis utiliser la méthode register_from_file de la classe UDFRegistration pour créer une UDF.

Voici des exemples d’utilisation de register_from_file.

Supposons que vous ayez un fichier Python test_udf_file.py qui contient :

def mod5(x: int) -> int:
    return x % 5

Ensuite, vous pouvez créer une UDF à partir de cette fonction du fichier test_udf_file.py.

# mod5() in that file has type hints
mod5_udf = session.udf.register_from_file(
      file_path="tests/resources/test_udf_dir/test_udf_file.py",
      func_name="mod5",
   )
   session.range(1, 8, 2).select(mod5_udf("id")).to_df("col1").collect()

[Row(COL1=1), Row(COL1=3), Row(COL1=0), Row(COL1=2)]

Vous pouvez également charger le fichier vers un emplacement de zone de préparation, puis l’utiliser pour créer l’UDF.

from snowflake.snowpark.types import IntegerType
# suppose you have uploaded test_udf_file.py to stage location @mystage.
mod5_udf = session.udf.register_from_file(
   file_path="@mystage/test_udf_file.py",
   func_name="mod5",
   return_type=IntegerType(),
   input_types=[IntegerType()],
)
session.range(1, 8, 2).select(mod5_udf("id")).to_df("col1").collect()

[Row(COL1=1), Row(COL1=3), Row(COL1=0), Row(COL1=2)]

Lecture de fichiers à partir d’une UDF¶

Pour lire le contenu d’un fichier, votre code Python peut :

• Lire un fichier spécifié de façon statique en important un fichier et en le lisant ensuite dans le répertoire personnel de l’UDF.

• Lire un fichier spécifié de façon dynamique avec SnowflakeFile. Vous pouvez le faire si vous avez besoin d’accéder à un fichier pendant un calcul.

import snowflake.snowpark as snowpark
from snowflake.snowpark.functions import udf
from snowflake.snowpark.files import SnowflakeFile
from snowflake.snowpark.types import StringType, IntegerType

@udf(name="get_file_length", replace=True, input_types=[StringType()], return_type=IntegerType(), packages=['snowflake-snowpark-python'])
def get_file_length(file_path):
  with SnowflakeFile.open(file_path) as f:
    s = f.read()
  return len(s);

session.sql("select get_file_length(build_scoped_file_url(@my_stage, 'example-file.txt'));")

Écriture de fichiers à partir de Python Snowpark UDFs et UDTFs¶

Avec Snowpark Python, vous pouvez écrire des fichiers dans des zones de préparation à l’aide de fonctions définies par l’utilisateur (UDFs), d’UDFs vectorisées, de fonctions de table définies par l’utilisateur (UDTFs) et d’UDTFs vectorisées. Dans le gestionnaire de la fonction, vous utilisez l’API SnowflakeFile pour ouvrir et écrire des fichiers. Lorsque vous renvoyez le fichier à partir de la fonction, le fichier est écrit en même temps que les résultats de la requête.

Une simple UDF pour écrire un fichier pourrait ressembler à ceci :

CREATE OR REPLACE FUNCTION write_file()
RETURNS STRING
LANGUAGE PYTHON
VOLATILE
RUNTIME_VERSION = 3.9
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'write_file'
AS
$$
from snowflake.snowpark.files import SnowflakeFile

def write_file():
  file = SnowflakeFile.open_new_result("w") # Open a new result file
  file.write("Hello world")                 # Write data
  return file               # File must be returned
$$;

L’exécution de cette UDF vous donnera alors une URL scopée faisant référence au fichier de résultat.

CREATE OR REPLACE FUNCTION convert_to_foo(filename string)
RETURNS STRING
LANGUAGE PYTHON
VOLATILE
RUNTIME_VERSION = 3.9
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'convert_to_foo'
AS
$$
from snowflake.snowpark.files import SnowflakeFile

def convert_to_foo(filename):
  input_file = SnowflakeFile.open(filename, "r")
  converted_file = SnowflakeFile.open_new_result("w")

  # Foo-type is just adding foo at the end of every line
  for line in input_file.readlines():
    converted_file.write(line[:-1] + 'foo' + '\n')
  return converted_file
$$;

COPY FILES INTO @output_stage FROM
  (SELECT convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, 'in.txt')), 'out.foo.txt');

CREATE TABLE files_to_convert(file string);
-- Populate files_to_convert with input files:
INSERT INTO files_to_convert VALUES ('file1.txt');
INSERT INTO files_to_convert VALUES ('file2.txt');

COPY FILES INTO @output_stage FROM
  (SELECT convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, file)),
      REPLACE(file, '.txt', '.foo.txt') FROM files_to_convert);

COPY FILES INTO @output_stage FROM
  (SELECT convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, RELATIVE_PATH)),
      REPLACE(RELATIVE_PATH, 'format1', 'format2') FROM DIRECTORY(@input_stage));

-- A basic UDF to read a file:
CREATE OR REPLACE FUNCTION read_udf(filename string)
RETURNS STRING
LANGUAGE PYTHON
VOLATILE
RUNTIME_VERSION = 3.9
PACKAGES = ('snowflake-snowpark-python')
HANDLER = 'read'
AS
'
from snowflake.snowpark.files import SnowflakeFile

def read(filename):
  return SnowflakeFile.open(filename, "r").read()
';

-- Create files_to_convert as in Example 2.

SELECT convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, file)) as new_file
  FROM files_to_convert;
-- The following query must be run within 24 hours from the prior one
SELECT read_udf(new_file) FROM TABLE(RESULT_SCAN(LAST_QUERY_ID()));

-- Set up files_to_convert as in Example 2.
-- Set up read_udf as in Example 4.

SELECT read_udf(
    convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, file))) FROM files_to_convert;

COPY FILES INTO @output_stage FROM
  (SELECT convert_to_foo(BUILD_SCOPED_FILE_URL(@input_stage, file)) FROM files_to_convert);

-- Refresh the directory to get new files in output_stage.
ALTER STAGE output_stage REFRESH;

-- Create a stage that includes the font (for PDF creation)

CREATE OR REPLACE STAGE fonts
URL = 's3://sfquickstarts/misc/';

-- UDF to write the data
CREATE OR REPLACE FUNCTION create_report_pdf(data string)
RETURNS TABLE (file string)
LANGUAGE PYTHON
RUNTIME_VERSION = 3.9
HANDLER='CreateReport'
PACKAGES = ('snowflake-snowpark-python', 'fpdf')
IMPORTS  = ('@fonts/DejaVuSans.ttf')
AS $$
from snowflake.snowpark.files import SnowflakeFile
from fpdf import FPDF
import shutil
import sys
import uuid
import_dir = sys._xoptions["snowflake_import_directory"]

class CreateReport:
  def __init__(self):
      self.pdf = None

  def process(self, data):
      if self.pdf == None:
        # PDF library edits this file, make sure it's unique.
        font_file = f'/tmp/DejaVuSans-{uuid.uuid4()}.ttf'
        shutil.copy(f'{import_dir}/DejaVuSans.ttf', font_file)
        self.pdf = FPDF()
        self.pdf.add_page()
        self.pdf.add_font('DejaVu', '', font_file, uni=True)
        self.pdf.set_font('DejaVu', '', 14)
      self.pdf.write(8, data)
      self.pdf.ln(8)

  def end_partition(self):
      f = SnowflakeFile.open_new_result("wb")
      f.write(self.pdf.output(dest='S').encode('latin-1'))
      yield f,
$$;

 -- Fictitious data
 CREATE OR REPLACE TABLE reportData(location string, item string);
 INSERT INTO reportData VALUES ('SanMateo' ,'Item A');
 INSERT INTO reportData VALUES ('SanMateo' ,'Item Z');
 INSERT INTO reportData VALUES ('SanMateo' ,'Item X');
 INSERT INTO reportData VALUES ('Bellevue' ,'Item B');
 INSERT INTO reportData VALUES ('Bellevue' ,'Item Q');

 -- Presigned URLs only work with SSE stages, see note above.
 CREATE OR REPLACE STAGE output_stage ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE');

 COPY FILES INTO @output_stage
   FROM (SELECT reports.file, location || '.pdf'
           FROM reportData, TABLE(create_report_pdf(item)
           OVER (partition BY location)) AS reports);

 -- Check the results
LIST @output_stage;
SELECT GET_PRESIGNED_URL(@output_stage, 'SanMateo.pdf');

CREATE OR REPLACE FUNCTION split_file(path string)
RETURNS TABLE(file string, name string)
LANGUAGE PYTHON
VOLATILE
PACKAGES = ('snowflake-snowpark-python')
RUNTIME_VERSION = 3.9
HANDLER = 'SplitCsvFile'
AS $$
import csv
from snowflake.snowpark.files import SnowflakeFile

class SplitCsvFile:
    def process(self, file):
      toTable1 = SnowflakeFile.open_new_result("w")
      toTable1Csv = csv.writer(toTable1)
      toTable2 = SnowflakeFile.open_new_result("w")
      toTable2Csv = csv.writer(toTable2)
      toTable3 = SnowflakeFile.open_new_result("w")
      toTable3Csv = csv.writer(toTable3)
      with SnowflakeFile.open(file, 'r') as file:
        # File is of the format 1:itemA \n 2:itemB \n [...]
        for line in file.readlines():
          forTable = line[0]
          if (forTable == "1"):
            toTable1Csv.writerow([line[2:-1]])
          if (forTable == "2"):
            toTable2Csv.writerow([line[2:-1]])
          if (forTable == "3"):
            toTable3Csv.writerow([line[2:-1]])
      yield toTable1, 'table1.csv'
      yield toTable2, 'table2.csv'
      yield toTable3, 'table3.csv'
$$;
-- Create a stage with access to an import file.
CREATE OR REPLACE STAGE staged_files url="s3://sfquickstarts/misc/";

-- Add the files to be split into a table - we just add one.
CREATE OR REPLACE TABLE filesToSplit(path string);
INSERT INTO filesToSplit VALUES ( 'items.txt');

-- Create output tables
CREATE OR REPLACE TABLE table1(item string);
CREATE OR REPLACE TABLE table2(item string);
CREATE OR REPLACE TABLE table3(item string);

-- Create output stage
CREATE OR REPLACE stage output_stage;

-- Creates files named path-tableX.csv
COPY FILES INTO @output_stage FROM
  (SELECT file, path || '-' || name FROM filesToSplit, TABLE(split_file(build_scoped_file_url(@staged_files, path))));

-- We use pattern and COPY INTO (not COPY FILES INTO) to upload to a final table.
COPY INTO table1 FROM @output_stage PATTERN = '.*.table1.csv';
COPY INTO table2 FROM @output_stage PATTERN = '.*.table2.csv';
COPY INTO table3 FROM @output_stage PATTERN = '.*.table3.csv';

-- See results
SELECT * from table1;
SELECT * from table2;
SELECT * from table3;

Utilisation des UDFs vectorisées¶

Les UDFs Python vectorisées vous permettent de définir des fonctions Python qui reçoivent des lots de lignes d’entrée sous forme de DataFrames Pandas et renvoient des lots de résultats sous forme de tableaux ou de séries Pandas. La colonne dans le dataframe Snowpark sera vectorisée comme une série Pandas à l’intérieur de l’UDF.

Voici un exemple d’utilisation de l’interface lots :

from sklearn.linear_model import LinearRegression
model = LinearRegression()
model.fit(X, y)

@udf(packages=['pandas', 'scikit-learn','xgboost'])
def predict(df: PandasDataFrame[float, float, float, float]) -> PandasSeries[float]:
    # The input pandas DataFrame doesn't include column names. Specify the column names explicitly when needed.
    df.columns = ["col1", "col2", "col3", "col4"]
    return model.predict(df)

Vous appelez des UDFs Python vectorisées de la même manière que vous appelez les autres UDFs Python. Pour plus d’informations, consultez UDFs Python vectorisées, qui explique comment créer une UDF vectorielle à l’aide d’une instruction SQL. Par exemple, vous pouvez utiliser le décorateur vectorized lorsque vous spécifiez le code Python dans l’instruction SQL. En utilisant l’API Snowpark Python décrite dans ce document, vous n’utilisez pas une instruction SQL pour créer une UDF vectorielle. On n’utilise donc pas le décorateur vectorized.

Il est possible de limiter le nombre de lignes par lot. Pour plus d’informations, voir Définir une taille de lot cible.

Pour plus d’explications et d’exemples sur l’utilisation de l’API Snowpark Python pour créer des UDFs vectorielles, reportez-vous à la section sur les UDFs de la référence API de Snowpark.