Funções definidas pelo usuário no Databricks Connect para Python

O Databricks Connect para Python dá suporte a UDF (funções definidas pelo usuário). Quando uma operação DataFrame que inclui UDFs é executada, as UDFs são serializadas pelo Databricks Connect e enviadas ao servidor como parte da solicitação.

Para obter informações sobre UDFs para Databricks Connect para Scala, consulte as funções definidas pelo usuário no Databricks Connect para Scala.

Definir um UDF

Para criar uma UDF no Databricks Connect para Python, use uma das seguintes funções com suporte:

• Funções definidas pelo usuário do PySpark
  • pyspark.sql.udf
  • pyspark.sql.pandas_udf
  • pyspark.sql.udtf
  • pyspark.sql.DataFrame.mapInPandas
  • pyspark.sql.DataFrame.mapInArrow
  • pyspark.sql.GroupedData.applyInPandas
  • pyspark.sql.GroupedData.applyInArrow
  • pyspark.sql.PandasCogroupedOps.applyInPandas
  • pyspark.sql.PandasCogroupedOps.applyInArrow
• Funções de streaming do PySpark
  • pyspark.sql.streaming.DataStreamWriter.foreach
  • pyspark.sql.streaming.DataStreamWriter.foreachBatch
  • pyspark.sql.streaming.StatefulProcessor

Por exemplo, o Python a seguir configura um UDF simples que ajusta os valores em uma coluna.

from pyspark.sql.functions import col, udf
from pyspark.sql.types import IntegerType
from databricks.connect import DatabricksSession

@udf(returnType=IntegerType())
def double(x):
    return x * x

spark = DatabricksSession.builder.getOrCreate()

df = spark.range(1, 2)
df = df.withColumn("doubled", double(col("id")))

df.show()

UDFs com dependências

O Databricks Connect dá suporte à especificação de dependências do Python necessárias para UDFs. Essas dependências são instaladas na computação do Databricks como parte do ambiente Python da UDF.

Esse recurso permite que os usuários especifiquem as dependências de que a UDF precisa além dos pacotes fornecidos no ambiente base. Ele também pode ser usado para instalar uma versão diferente do pacote do que é fornecido no ambiente base.

As dependências podem ser instaladas a partir das seguintes fontes:

• Pacotes de PyPI
  • Os pacotes PyPI podem ser especificados de acordo com o PEP 508, por exemplo, dicepyjokes<1 ou simplejson==3.19.*.
• Arquivos armazenados em volumes do Catálogo do Unity
  • Há suporte para pacotes wheel (.whl) e arquivos tar compactados com gzip (.tar.gz). O usuário deve receber permissão READ_FILE no arquivo do volume re:[UC].
  • Ao instalar pacotes a partir de volumes do Unity Catalog, para invocar as UDFs, os usuários precisam de READ VOLUME permissão no volume de origem. Conceder essa permissão a todos os usuários da conta habilita automaticamente isso para novos usuários.
  • Os arquivos de volumes do Catálogo do Unity devem ser especificados como dbfs:<path>, por exemplo, dbfs:/Volumes/users/someone@example.com/wheels/my_private_dep.whl ou dbfs:/Volumes/users/someone@example.com/tars/my_private_deps.tar.gz.

Para incluir dependências personalizadas em seu UDF, especifique-as em um ambiente usando withDependenciese use esse ambiente para criar uma sessão do Spark. As dependências são instaladas em sua computação do Databricks e estarão disponíveis em todos os UDFs que usam essa sessão do Spark.

O código a seguir declara o pacote dice PyPI como uma dependência:

from databricks.connect import DatabricksSession, DatabricksEnv
env = DatabricksEnv().withDependencies("dice==3.1.0")
spark = DatabricksSession.builder.withEnvironment(env).getOrCreate()

Ou, para especificar uma dependência de uma roda em um volume:

from databricks.connect import DatabricksSession, DatabricksEnv

env = DatabricksEnv().withDependencies("dbfs:/Volumes/users/someone@example.com/wheels/my_private_dep.whl")
spark = DatabricksSession.builder.withEnvironment(env).getOrCreate()

Comportamento em notebooks e trabalhos do Databricks

Em notebooks e trabalhos, as dependências de UDF precisam ser instaladas diretamente no REPL. O Databricks Connect valida o ambiente do Python REPL verificando se todas as dependências especificadas já estão instaladas e gera uma exceção se alguma não estiver instalada.

A validação do ambiente do notebook ocorre tanto para as dependências de volume do PyPI quanto para as do Unity Catalog. As dependências de volume precisam ser empacotadas seguindo as especificações de empacotamento padrão do Python do PEP-427 ou posterior para arquivos wheel e PEP-241 ou posterior para arquivos de distribuição de origem. Para obter mais informações sobre os padrões de empacotamento do Python, consulte a documentação do PyPA.

Limitações

• Arquivos como a roda do Python ou a distribuição de origem em seu computador de desenvolvimento local não podem ser especificados diretamente como uma dependência. Eles devem primeiro ser carregados nos volumes do Catálogo do Unity.
• As dependências UDF requerem suporte para pyspark.sql.streaming.DataStreamWriter.foreach e pyspark.sql.streaming.DataStreamWriter.foreachBatch, o que exige o Databricks Connect para Python na versão 18.0 ou superior, assim como um cluster executando o Databricks Runtime na versão 18.0 ou superior.
• Dependências UDF não são suportadas para UDFs de agregação de pandas em funções de janela.

Exemplos

O exemplo a seguir define dependências de PyPI e volumes em um ambiente, cria uma sessão com esse ambiente e, em seguida, define e chama UDFs que usam essas dependências:

from databricks.connect import DatabricksSession, DatabricksEnv
from pyspark.sql.functions import udf, col, pandas_udf
from pyspark.sql.types import IntegerType, LongType, StringType
import pandas as pd

pypi_deps = ["pyjokes>=0.8,<1"]

volumes_deps = [
    # Example library from: https://pypi.org/project/dice/#files
    "dbfs:/Volumes/main/someone@example.com/test/dice-4.0.0-py3-none-any.whl"
    # Example library from: https://pypi.org/project/simplejson/#files
    "dbfs:/Volumes/main/someone@example.com/test/simplejson-3.19.3.tar.gz",
]

env = DatabricksEnv().withDependencies(pypi_deps).withDependencies(volumes_deps)
spark = DatabricksSession.builder.withEnvironment(env).getOrCreate()

# UDFs
@udf(returnType=StringType())
def get_joke():
    from pyjokes import get_joke
    return get_joke()

@udf(returnType=IntegerType())
def double_and_json_parse(x):
    import simplejson
    return simplejson.loads(simplejson.dumps(x * 2))


@pandas_udf(returnType=LongType())
def multiply_and_add_roll(a: pd.Series, b: pd.Series) -> pd.Series:
    import dice
    return a * b + dice.roll(f"1d10")[0]


df = spark.range(1, 10)
df = df.withColumns({
    "joke": get_joke(),
    "doubled": double_and_json_parse(col("id")),
    "mutliplied_with_roll": multiply_and_add_roll(col("id"), col("doubled"))
})
df.show()

Ambiente base do Python

Os UDFs são executados na computação do Databricks e não no cliente. O ambiente base do Python no qual as UDFs são executadas depende da computação do Databricks.

Para clusters, o ambiente base do Python é o ambiente Python da versão do Databricks Runtime em execução no cluster. A versão do Python e a lista de pacotes nesse ambiente base são encontradas no ambiente do sistema e nas seções de bibliotecas do Python instaladas das notas de versão do Databricks Runtime.

Para computação sem servidor, o ambiente python base corresponde à versão do ambiente sem servidor de acordo com a tabela a seguir.