Fonctions de table définies par l'utilisateur (UDTF) Python
Aperçu
Cette fonctionnalité est en Aperçu public dans Databricks Runtime 14.3 LTS et versions supérieures.
Une fonction de table définie par l'utilisateur (UDTF) vous permet d'enregistrer des fonctions qui renvoient des tables au lieu de valeurs scalaires. Contrairement aux fonctions scalaires qui renvoient une valeur de résultat unique à chaque appel, chaque UDTF est appelée dans la clause FROM d'une instruction SQL et renvoie une table entière en sortie.
Chaque appel UDTF peut accepter zéro ou plusieurs arguments. Ces arguments peuvent être des expressions scalaires ou des arguments de table représentant des tables d'entrée entières.
Les UDTF peuvent être enregistrées de deux manières :
- Unity Catalog : Enregistrez l'UDTF en tant qu'objet gouverné dans Unity Catalog. Consultez les fonctions de table définies par l'utilisateur (UDTF) Python dans Unity Catalog.
- Limitée à la session : S'enregistrer auprès du
SparkSessionlocal, isolé au Notebook ou Job actuel.
Databricks recommande d'enregistrer les fonctions UDTF dans Unity Catalog pour profiter d'une gouvernance centralisée qui facilite le partage sécurisé et la réutilisation des fonctions entre les utilisateurs et les équipes.
Syntaxe UDTF de base
Apache Spark implémente les UDTF Python comme des classes Python avec une méthode eval obligatoire qui utilise yield pour émettre des lignes de sortie.
Pour utiliser votre classe en tant que UDTF, vous devez importer la fonction PySpark udtf. Databricks recommande
d'utiliser cette fonction comme décorateur et de spécifier explicitement les noms et types de champs à l'aide de l'option
returnType (à moins que la classe ne définisse une méthode analyze comme décrit dans une section ultérieure).
La fonction UDTF suivante crée une table en utilisant une liste fixe de deux arguments entiers :
from pyspark.sql.functions import lit, udtf
@udtf(returnType="sum: int, diff: int")
class GetSumDiff:
def eval(self, x: int, y: int):
yield x + y, x - y
GetSumDiff(lit(1), lit(2)).show()
+----+-----+
| sum| diff|
+----+-----+
| 3| -1|
+----+-----+
Enregistrer une UDTF
Pour enregistrer une UDTF à l'endu00e9e de la session pour une utilisation dans les queries SQL, utilisez spark.udtf.register(). Fournissez un nom pour la fonction SQL et la classe UDTF Python.
spark.udtf.register("get_sum_diff", GetSumDiff)
Appeler une fonction de table définie par l'utilisateur (UDTF) enregistrée
Une fois enregistré, vous pouvez utiliser l'UDTF dans SQL en utilisant soit la commande magique %sql ou la fonction spark.sql()
:
spark.udtf.register("get_sum_diff", GetSumDiff)
spark.sql("SELECT * FROM get_sum_diff(1,2);").show()
%sql
SELECT * FROM get_sum_diff(1,2);
Mettre à niveau une UDTF à portée de session vers Unity Catalog
Aperçu
L'enregistrement des fonctions UDTF Python dans Unity Catalog est en aperçu public. Les UDTF d'Unity Catalog nécessitent Databricks Runtime version 17.1 ou une version supérieure. Voir Exigences.
Vous pouvez mettre à niveau une UDTF à portée de session vers Unity Catalog pour bénéficier de la gouvernance centralisée et faciliter le partage et la réutilisation sécurisés des fonctions entre les utilisateurs et les équipes.
Pour mettre à niveau une UDTF délimitée à la session vers Unity Catalog, utilisez SQL DDL avec l’instruction CREATE OR REPLACE FUNCTION. L’exemple suivant montre comment convertir l’UDTF GetSumDiff d’une fonction délimitée à la session en une fonction Unity Catalog :
CREATE OR REPLACE FUNCTION get_sum_diff(x INT, y INT)
RETURNS TABLE (sum INT, diff INT)
LANGUAGE PYTHON
HANDLER 'GetSumDiff'
AS $$
class GetSumDiff:
def eval(self, x: int, y: int):
yield x + y, x - y
$$;
SELECT * FROM get_sum_diff(10, 3);
+-----+------+
| sum | diff |
+-----+------+
| 13 | 7 |
+-----+------+
Pour plus d'informations sur les UDTF d'Unity Catalog, consultez les fonctions de table définies par l'utilisateur Python (UDTF) dans Unity Catalog.
Utiliser Apache Arrow
Si votre UDTF reçoit une petite quantité de données en entrée, mais génère une grande table en sortie, Databricks
vous recommande d’utiliser Apache Arrow. Vous pouvez l'activer en spécifiant le paramètre useArrow lors de la
déclaration de la fonction UDTF :
@udtf(returnType="c1: int, c2: int", useArrow=True)
Listes d'arguments variables — *args et **kwargs
Vous pouvez utiliser la syntaxe Python *args ou **kwargs et implémenter une logique pour gérer un nombre non spécifié de valeurs d'entrée.
L’exemple suivant renvoie le même résultat tout en vérifiant explicitement la longueur d’entrée et les types des arguments :
@udtf(returnType="sum: int, diff: int")
class GetSumDiff:
def eval(self, *args):
assert(len(args) == 2)
assert(isinstance(arg, int) for arg in args)
x = args[0]
y = args[1]
yield x + y, x - y
GetSumDiff(lit(1), lit(2)).show()
Voici le même exemple, mais en utilisant des arguments du mot clé :
@udtf(returnType="sum: int, diff: int")
class GetSumDiff:
def eval(self, **kwargs):
x = kwargs["x"]
y = kwargs["y"]
yield x + y, x - y
GetSumDiff(x=lit(1), y=lit(2)).show()
Définir un schéma statique au moment de l'enregistrement
La fonction UDTF renvoie des lignes avec un schéma de sortie comprenant une séquence ordonnée de noms et de types de colonnes. Si le schéma UDTF doit toujours rester le même pour toutes les query, vous pouvez spécifier un schéma statique fixe après le décorateur @udtf. Il doit s’agir d’un StructType:
StructType().add("c1", StringType())
Ou une chaîne DDL représentant un type de structure :
c1: string
Calculer un schéma dynamique au moment de l'appel de fonction
Les UDTF peuvent également calculer le schéma de sortie par programmation pour chaque appel en fonction des valeurs des arguments d'entrée. Pour ce faire, définissez une méthode statique appelée analyze qui accepte zéro ou plusieurs parameters qui correspondent aux arguments fournis à l'appel UDTF spécifique.
Chaque argument de la méthode analyze est une instance de la classe AnalyzeArgument qui contient les champs suivants :
| Description |
|---|---|
| Le type de l’argument d’entrée en tant que |
| La valeur de l'argument d'entrée en tant que |
| Indique si l'argument d'entrée est une table en tant que |
| Si l'argument d'entrée est une expression à réduction constante en tant que |
La méthode analyze renvoie une instance de la classe AnalyzeResult, qui inclut le schéma de la table de résultats en tant que StructType, ainsi que des champs facultatifs. Si l'UDTF accepte un argument de table d'entrée, alors le AnalyzeResult peut également inclure une méthode demandée pour partitionner et ordonner les lignes de la table d'entrée sur plusieurs appels UDTF, comme décrit ci-après.
| Description |
|---|---|
| Le schéma de la table de résultats sous la forme d'un |
| Envoyer toutes les lignes d'entrée à la même instance de classe UDTF comme un |
| S'il est défini comme non vide, toutes les lignes ayant chaque combinaison unique de valeurs des expressions de partitionnement sont consommées par une instance distincte de la classe UDTF. |
| S’il est défini comme non vide, il spécifie un ordre des lignes au sein de chaque partition. |
| S'il est défini comme non vide, il s'agit d'une séquence d'expressions que la fonction UDTF spécifie pour que Catalyst évalue les colonnes de l'argument TABLE d'entrée. La fonction UDTF reçoit un attribut d'entrée pour chaque nom de la liste dans l'ordre où ils sont répertoriés. |
Cet exemple analyze renvoie une colonne de sortie pour chaque mot de l’argument de chaîne d’entrée.
from pyspark.sql.functions import lit, udtf
from pyspark.sql.types import StructType, IntegerType
from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult
@udtf
class MyUDTF:
@staticmethod
def analyze(text: AnalyzeArgument) -> AnalyzeResult:
schema = StructType()
for index, word in enumerate(sorted(list(set(text.value.split(" "))))):
schema = schema.add(f"word_{index}", IntegerType())
return AnalyzeResult(schema=schema)
def eval(self, text: str):
counts = {}
for word in text.split(" "):
if word not in counts:
counts[word] = 0
counts[word] += 1
result = []
for word in sorted(list(set(text.split(" ")))):
result.append(counts[word])
yield result
MyUDTF(lit("hello world")).columns
['word_0', 'word_1']
Transférer l'état aux futurs appels de eval
La méthode analyze peut servir d'emplacement pratique pour effectuer l'initialisation, puis transférer les résultats aux futures invocations de la méthode eval pour le même appel UDTF.
Pour ce faire, créez une sous-classe de AnalyzeResult et renvoyez une instance de la sous-classe à partir de la méthode analyze.
Ensuite, ajoutez un argument supplémentaire à la méthode __init__ pour accepter cette instance.
Cet exemple analyze renvoie un schéma de sortie constant, mais ajoute des informations personnalisées dans les métadonnées de résultat pour être consommées par les futurs appels de méthode __init__ :
from pyspark.sql.functions import lit, udtf
from pyspark.sql.types import StructType, IntegerType
from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult
@dataclass
class AnalyzeResultWithBuffer(AnalyzeResult):
buffer: str = ""
@udtf
class TestUDTF:
def __init__(self, analyze_result=None):
self._total = 0
if analyze_result is not None:
self._buffer = analyze_result.buffer
else:
self._buffer = ""
@staticmethod
def analyze(argument, _) -> AnalyzeResult:
if (
argument.value is None
or argument.isTable
or not isinstance(argument.value, str)
or len(argument.value) == 0
):
raise Exception("The first argument must be a non-empty string")
assert argument.dataType == StringType()
assert not argument.isTable
return AnalyzeResultWithBuffer(
schema=StructType()
.add("total", IntegerType())
.add("buffer", StringType()),
withSinglePartition=True,
buffer=argument.value,
)
def eval(self, argument, row: Row):
self._total += 1
def terminate(self):
yield self._total, self._buffer
spark.udtf.register("test_udtf", TestUDTF)
spark.sql(
"""
WITH t AS (
SELECT id FROM range(1, 21)
)
SELECT total, buffer
FROM test_udtf("abc", TABLE(t))
"""
).show()
+-------+-------+
| count | buffer|
+-------+-------+
| 20 | "abc"|
+-------+-------+
Générer des lignes de sortie
La méthode eval s'exécute une fois pour chaque ligne de l'argument de table d'entrée (ou une seule fois si aucun argument de table n'est fourni), suivie d'une invocation de la méthode terminate à la fin. L'une ou l'autre méthode produit zéro ou plusieurs lignes qui sont conformes au schéma de résultats en renvoyant des tuples, des listes ou des objets pyspark.sql.Row.
Cet exemple renvoie une ligne en fournissant un tuple de trois éléments :
def eval(self, x, y, z):
yield (x, y, z)
Vous pouvez également omettre les parenthèses :
def eval(self, x, y, z):
yield x, y, z
Ajoutez une virgule de fin pour renvoyer une ligne avec une seule colonne :
def eval(self, x, y, z):
yield x,
Vous pouvez également produire un objet pyspark.sql.Row.
def eval(self, x, y, z):
from pyspark.sql.types import Row
yield Row(x, y, z)
Cet exemple génère des lignes de sortie à partir de la méthode terminate à l'aide d'une liste Python. Vous pouvez stocker l'état dans la classe à partir des étapes précédentes de l'évaluation UDTF à cette fin.
def terminate(self):
yield [self.x, self.y, self.z]
Transmettre des arguments scalaires à une UDTF
Vous pouvez passer des arguments scalaires à une UDTF en tant qu'expressions constantes comprenant des valeurs littérales ou des fonctions basées sur celles-ci. Par exemple :
SELECT * FROM get_sum_diff(1, y => 2)
Transmettre les arguments de table à une UDTF
Les UDTF Python peuvent accepter une table d'entrée comme argument en plus des arguments d'entrée scalaires. Une seule UDTF peut également accepter un argument de table et plusieurs arguments scalaires.
Ensuite, toute query SQL peut fournir une table d'entrée à l'aide du mot-clé TABLE suivi de parenthèses
entourant un identifiant de table approprié, comme TABLE(t). Vous pouvez également transmettre une sous-requête de table, comme TABLE(SELECT a, b, c FROM t) ou TABLE(SELECT t1.a, t2.b FROM t1 INNER JOIN t2 USING (key)).
L'argument de table d'entrée est ensuite représenté comme un argument pyspark.sql.Row de la méthode eval, avec un appel à la méthode eval pour chaque ligne de la table d'entrée. Vous pouvez utiliser les annotations de champ de colonne PySpark standard pour interagir avec les colonnes de chaque ligne. L'exemple suivant montre l'importation explicite du type PySpark Row, puis le filtrage de la table transmise sur le champ id :
from pyspark.sql.functions import udtf
from pyspark.sql.types import Row
@udtf(returnType="id: int")
class FilterUDTF:
def eval(self, row: Row):
if row["id"] > 5:
yield row["id"],
spark.udtf.register("filter_udtf", FilterUDTF)
Pour interroger la fonction, utilisez le mot-clé SQL TABLE :
SELECT * FROM filter_udtf(TABLE(SELECT * FROM range(10)));
+---+
| id|
+---+
| 6|
| 7|
| 8|
| 9|
+---+
Spécifier un partitionnement des lignes d'entrée issues des appels de fonction
Lorsque vous appelez une UDTF avec un argument de table, toute query SQL peut partitionner la table d'entrée sur plusieurs appels UDTF en fonction des valeurs d'une ou plusieurs colonnes de la table d'entrée.
Pour spécifier une partition, utilisez la clause PARTITION BY dans l'appel de fonction après l'argument TABLE.
Ceci garantit que toutes les lignes d'entrée avec chaque combinaison unique de valeurs des colonnes de partitionnement seront consommées par une seule instance de la classe UDTF.
Notez qu'en plus des références de colonne simples, la clause PARTITION BY accepte également des expressions arbitraires basées sur les colonnes de la table d'entrée. Par exemple, vous pouvez spécifier le LENGTH d'une chaîne de caractères, extraire un mois d'une date ou concaténer deux valeurs.
Il est également possible de spécifier WITH SINGLE PARTITION au lieu de PARTITION BY pour ne demander qu'une seule partition dans laquelle toutes les lignes d'entrée doivent être consommées par exactement une instance de la classe UDTF.
Dans chaque partition, vous pouvez éventuellement spécifier un ordre requis des lignes d'entrée lorsque la méthode eval de l'UDTF les consomme. Pour ce faire, fournissez une clause ORDER BY après la clause PARTITION BY ou WITH SINGLE PARTITION décrite ci-dessus.
Par exemple, considérez la fonction UDTF suivante :
from pyspark.sql.functions import udtf
from pyspark.sql.types import Row
@udtf(returnType="a: string, b: int")
class FilterUDTF:
def __init__(self):
self.key = ""
self.max = 0
def eval(self, row: Row):
self.key = row["a"]
self.max = max(self.max, row["b"])
def terminate(self):
yield self.key, self.max
spark.udtf.register("filter_udtf", FilterUDTF)
Vous pouvez spécifier des options de partitionnement lors de l'appel de l'UDTF sur la table d'entrée de plusieurs façons :
-- Create an input table with some example values.
DROP TABLE IF EXISTS values_table;
CREATE TABLE values_table (a STRING, b INT);
INSERT INTO values_table VALUES ('abc', 2), ('abc', 4), ('def', 6), ('def', 8);
SELECT * FROM values_table;
+-------+----+
| a | b |
+-------+----+
| "abc" | 2 |
| "abc" | 4 |
| "def" | 6 |
| "def" | 8 |
+-------+----+
-- Query the UDTF with the input table as an argument and a directive to partition the input
-- rows such that all rows with each unique value in the `a` column are processed by the same
-- instance of the UDTF class. Within each partition, the rows are ordered by the `b` column.
SELECT * FROM filter_udtf(TABLE(values_table) PARTITION BY a ORDER BY b) ORDER BY 1;
+-------+----+
| a | b |
+-------+----+
| "abc" | 4 |
| "def" | 8 |
+-------+----+
-- Query the UDTF with the input table as an argument and a directive to partition the input
-- rows such that all rows with each unique result of evaluating the "LENGTH(a)" expression are
-- processed by the same instance of the UDTF class. Within each partition, the rows are ordered
-- by the `b` column.
SELECT * FROM filter_udtf(TABLE(values_table) PARTITION BY LENGTH(a) ORDER BY b) ORDER BY 1;
+-------+---+
| a | b |
+-------+---+
| "def" | 8 |
+-------+---+
-- Query the UDTF with the input table as an argument and a directive to consider all the input
-- rows in one single partition such that exactly one instance of the UDTF class consumes all of
-- the input rows. Within each partition, the rows are ordered by the `b` column.
SELECT * FROM filter_udtf(TABLE(values_table) WITH SINGLE PARTITION ORDER BY b) ORDER BY 1;
+-------+----+
| a | b |
+-------+----+
| "def" | 8 |
+-------+----+
Spécifiez un partitionnement des lignes d'entrée de la méthode analyze
Veuillez noter que pour chacune des méthodes de partitionnement de la table d'entrée mentionnées ci-dessus lors de l'appel de fonctions UDTF dans les requêtes SQL, il existe une méthode correspondante pour la méthode analyze de la fonction UDTF afin de spécifier automatiquement la même méthode de partitionnement à la place.
- Au lieu d’appeler une UDTF comme
SELECT * FROM udtf(TABLE(t) PARTITION BY a), vous pouvez mettre à jour la méthodeanalyzepour définir le champpartitionBy=[PartitioningColumn("a")]et appeler simplement la fonction à l’aide deSELECT * FROM udtf(TABLE(t)). - De même, au lieu de spécifier
TABLE(t) WITH SINGLE PARTITION ORDER BY bdans la query SQL, vous pouvez faire en sorte queanalyzedéfinisse les champswithSinglePartition=trueetorderBy=[OrderingColumn("b")], puis passer simplementTABLE(t). - Au lieu de passer
TABLE(SELECT a FROM t)dans la query SQL, vous pouvez faire en sorte queanalyzedéfinisseselect=[SelectedColumn("a")]et ensuite simplement passerTABLE(t).
Dans l'exemple suivant, analyze renvoie un schéma de sortie constant, sélectionne un sous-ensemble de colonnes de la table d'entrée et spécifie que la table d'entrée est partitionnée sur plusieurs appels UDTF en fonction des valeurs de la colonne date :
@staticmethod
def analyze(*args) -> AnalyzeResult:
"""
The input table will be partitioned across several UDTF calls based on the monthly
values of each `date` column. The rows within each partition will arrive ordered by the `date`
column. The UDTF will only receive the `date` and `word` columns from the input table.
"""
from pyspark.sql.functions import (
AnalyzeResult,
OrderingColumn,
PartitioningColumn,
)
assert len(args) == 1, "This function accepts one argument only"
assert args[0].isTable, "Only table arguments are supported"
return AnalyzeResult(
schema=StructType()
.add("month", DateType())
.add("longest_word", IntegerType()),
partitionBy=[
PartitioningColumn("extract(month from date)")],
orderBy=[
OrderingColumn("date")],
select=[
SelectedColumn("date"),
SelectedColumn(
name="length(word)",
alias="length_word")])