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

O Databricks Connect for Python suporta funções definidas pelo usuário (UDF). 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 for Scala, consulte Funções definidas pelo usuário no Databricks Connect for Scala.

Definir uma UDF

Para criar um UDF no Databricks Connect for Python, use uma das seguintes funções suportadas:

• 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 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 quadra 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 suporta a especificação de dependências Python que são necessárias para UDFs. Essas dependências são instaladas na plataforma Databricks como parte do ambiente Python da UDF.

Esse recurso permite que os usuários especifiquem dependências que o 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 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 Unity
  • Ambos os arquivos wheel (.whl) e os arquivos tar compactados (.tar.gz) são suportados. O utilizador deve receber permissão READ_FILE no ficheiro no volume re:[UC].
  • Ao instalar pacotes a partir de volumes do Catálogo Unity, para invocar os 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 isso automaticamente para novos usuários.
  • Os arquivos de volumes do Catálogo 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 withDependencieso e, em seguida, use esse ambiente para criar uma sessão do Spark. As dependências são instaladas na sua infraestrutura computacional Databricks e estarão disponíveis em todas as UDFs que utilizam esta 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 tarefas Databricks

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

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

Limitações

• Arquivos como roda Python ou distribuição de código-fonte em sua máquina de desenvolvimento local não podem ser especificados diretamente como uma dependência. Eles devem primeiro ser carregados nos volumes do Catálogo Unity.
• O suporte de dependências UDF para pyspark.sql.streaming.DataStreamWriter.foreach e pyspark.sql.streaming.DataStreamWriter.foreachBatch requer Databricks Connect para Python 18.0 ou superior, e um cluster a correr Databricks Runtime 18.0 ou superior.
• As dependências de UDF não são suportadas para UDFs de agregação de pandas sobre 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 Python

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

Para clusters, o ambiente Python base é o ambiente Python da versão do Databricks Runtime em execução no cluster. A versão do Python e a lista de pacotes neste ambiente base são encontradas nas seções Ambiente do sistema e Bibliotecas 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.