SDK do Databricks para Python

Observação

Databricks recomenda o Databricks ativo Bundles para criar, desenvolver, implantar e testar o Job e outros Databricks recursos como código-fonte. Veja o que são Databricks ativo Bundles?

Neste artigo, o senhor aprenderá como automatizar Databricks operações e acelerar o desenvolvimento com o Databricks SDK para Python. Este artigo complementa Databricks SDK Python a Databricks SDK Python documentação para em Read The Docs e os exemplos de código no para repositório GitHub em.

Observação

O SDK da Databricks para Python está em versão beta e pode ser usado em produção.

Durante o período Beta, o Databricks recomenda marcar uma dependência na versão secundária específica do SDK do Databricks para Python da qual seu código depende. Por exemplo, você pode marcar dependências em arquivos como requirements.txt para venv ou pyproject.toml e poetry.lock para Poetry. Para mais informações sobre marcação de dependências, consulte Ambientes virtuais e pacote para venv ou Instalação de dependências para Poetry.

Antes de começar

Você pode usar o SDK do Databricks para Python de dentro de um notebook do Databricks ou de sua máquina de desenvolvimento local.

Antes de começar a usar o SDK do Databricks para Python, sua máquina de desenvolvimento deve ter:

  • Autenticação do Databricks configurada.

  • Python 3.8 ou superior instalado. Para automatizar os recursos de compute do Databricks, você deve ter as versões principais e secundárias do Python que correspondam à versão instalada no recurso de compute do Databricks de destino. Os exemplos deste artigo se baseiam na automação de clusters com o Databricks Runtime 13.3 LTS, que tem o Python 3.10 instalado. Para a versão correta, consulte as notas sobre a versão e compatibilidade do Databricks Runtime para a versão do Databricks Runtime do seu cluster.

  • A Databricks recomenda que o senhor crie e ative um ambiente virtual Python para cada projeto Python que usar com o Databricks SDK for Python. Python Os ambientes virtuais ajudam a garantir que seu projeto de código esteja usando versões compatíveis de Python e Python pacote (neste caso, o Databricks SDK para Python pacote). Para obter mais informações sobre os ambientes virtuais do Python, consulte venv ou Poetry.

Introdução ao SDK do Databricks para Python

Esta seção descreve como começar a usar o SDK do Databricks para Python em sua máquina de desenvolvimento local. Para usar o SDK do Databricks para Python de dentro de um notebook do Databricks, vá para Usar o SDK do Databricks para Python a partir de um notebook do Databricks.

  1. Em sua máquina de desenvolvimento, com a autenticação do Databricks configurada, o Python já instalado e o ambiente virtual do Python já ativo, instale o pacote databricks-sdk (e suas dependências) a partir do Python Package Index (PyPI), como segue:

    Use pip para instalar o pacote databricks-sdk . (Em alguns sistemas, talvez o senhor precise substituir pip3 por pip, aqui e em toda parte).

    pip3 install databricks-sdk

    poetry add databricks-sdk

    Para instalar uma versão específica do pacote databricks-sdk enquanto o SDK do Databricks para Python estiver na versão beta, consulte o histórico de lançamento do pacote. Por exemplo, para instalar a versão 0.1.6:

    pip3 install databricks-sdk==0.1.6

    poetry add databricks-sdk==0.1.6

    Para atualizar uma instalação existente do pacote Databricks SDK for Python para a versão mais recente, execute o seguinte comando:

    pip3 install --upgrade databricks-sdk

    poetry add databricks-sdk@latest

    Para mostrar o Version atual do pacote do SDK do Databricks para Python e outros detalhes, execute o seguinte comando:

    pip3 show databricks-sdk

    poetry show databricks-sdk

  2. Em seu ambiente virtual do Python, crie um arquivo de código do Python que importa o SDK do Databricks para Python. O exemplo a seguir, em um arquivo denominado main.py com o seguinte conteúdo, simplesmente lista todos os clusters no seu workspace do Databricks:

    from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

for c in w.clusters.list():
  print(c.cluster_name)

  3. Execute seu arquivo de código Python, supondo que seja um arquivo denominado main.py, executando o comando python:

    python3.10 main.py

    Se você estiver no shell do ambiente virtual:

    python3.10 main.py

    Se você não estiver no shell do ambiente virtual:

    poetry run python3.10 main.py

    Observação

    Ao não definir nenhum argumento na chamada anterior como w = WorkspaceClient(), o Databricks SDK para Python usa seu processo padrão para tentar executar a autenticação do Databricks. Para substituir este comportamento padrão, consulte a seção de autenticação.

Autentique o SDK do Databricks para Python com sua conta ou workspace do Databricks

Esta seção descreve como autenticar o SDK do Databricks para Python a partir de sua máquina de desenvolvimento local para sua conta ou workspace do Databricks. Para autenticar o SDK do Databricks para Python de dentro de um notebook do Databricks, vá para Usar o SDK do Databricks para Python a partir de um notebook do Databricks.

O SDK do Databricks para Python implementa o padrão de autenticação unificada do cliente Databricks, uma abordagem arquitetônica e programática consolidada e consistente para autenticação. Essa abordagem ajuda a tornar a configuração e a automação da autenticação com o Databricks mais centralizada e previsível. Ela permite configurar a autenticação do Databricks uma vez e, em seguida, usar essa configuração em várias ferramentas e SDKs do Databricks sem mais alterações na configuração de autenticação. Para mais informações, incluindo exemplos de código mais completos no Python, consulte Autenticação unificada do cliente do Databricks.

Alguns dos padrões de codificação disponíveis para inicializar a autenticação do Databricks com o SDK do Databricks para Python incluem:

  • Use a autenticação padrão do Databricks seguindo um destes procedimentos:

    • Crie ou identifique um perfil de configuração do Databricks personalizado com os campos obrigatórios para o tipo de autenticação desejada do Databricks. Em seguida, defina a variável de ambiente DATABRICKS_CONFIG_PROFILE com o nome do perfil de configuração personalizado.

    • Defina as variáveis de ambiente necessárias para o tipo de autenticação do Databricks de destino.

    Em seguida, instancie, por exemplo, um objeto WorkspaceClient com autenticação default do Databricks da seguinte maneira:

    from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
# ...

  • Embora seja possível codificar diretamente os campos necessários, não é aconselhável, pois isso pode expor informações confidenciais em seu código, como access tokens pessoais do Databricks. O exemplo a seguir codifica o host do Databricks e valores de access token para autenticação de token do Databricks:

    from databricks.sdk import WorkspaceClient

w = WorkspaceClient(
  host  = 'https://...',
  token = '...'
)
# ...

Veja também Autenticação na documentação do Databricks SDK para Python.

Usar o SDK do Databricks para Python a partir de um notebook do Databricks

O senhor pode chamar Databricks SDK para a funcionalidade Python de um notebook Databricks que tenha um Databricks cluster anexado com o Databricks SDK para Python instalado. Ele é instalado por default em todos os Databricks clusters que usam Databricks Runtime 13.3 LTS ou acima. Para Databricks clusters que usam Databricks Runtime 12.2 LTS e abaixo, o senhor deve instalar primeiro o Databricks SDK para Python. Consulte o passo 1: Instalar ou atualizar o Databricks SDK para Python.

Para ver a versão Databricks SDK para Python que está instalada para uma versão específica Databricks Runtime, consulte a seção Installed Python biblioteca das Databricks Runtime notas sobre a versão para essa versão.

Databricks recomenda que o senhor instale a versão mais recente disponível do SDK do PiPy, mas, no mínimo, instale ou atualize para Databricks SDK para Python 0.6.0 ou acima, pois default Databricks A autenticação do Notebook é usada pela versão 0.6.0 e pelo acima em todas as versões Databricks Runtime.

Observação

Databricks Runtime O 15.1 é o primeiro Databricks Runtime a ter uma versão do Databricks SDK para Python (0.20.0) instalada que oferece suporte à autenticação do default Notebook sem necessidade de atualização.

A tabela a seguir descreve o suporte à autenticação do Notebook para as versões Databricks SDK para Python e Databricks Runtime:

SDK/DBR

10.4 LTS

11.3 LTS

12.3 LTS

13.3 LTS

14.3 LTS

15.1 e acima

0.1.7 e abaixo

0.1.10

✓

✓

✓

✓

✓

0.6.0

✓

✓

✓

✓

✓

✓

0,20,0 e acima

✓

✓

✓

✓

✓

✓

A autenticação default do notebook do Databricks depende de um access token pessoal temporário do Databricks que o Databricks gera automaticamente em segundo plano para seu próprio uso. O Databricks exclui esse token temporário depois que o notebook para de funcionar.

Importante

  • A autenticação padrão do notebook Databricks funciona somente no nó do driver do cluster e não em nenhum dos nós worker ou executor do cluster.

  • A autenticação do notebook do Databricks não funciona com perfis de configuração do Databricks.

Se Databricks accounto senhor quiser chamar -level APIs ou se quiser usar um Databricks tipo de autenticação que não seja default Databricks Notebook authentication, também há suporte para os seguintes tipos de autenticação:

Tipo de autenticação

Versões do SDK do Databricks para Python

Autenticação OAuth machine-to-machine (M2M)

0.19.0 e superior

Autenticação OAuth user-to-machine (U2M)

0.19.0 e superior

Autenticação de credenciais do Google Cloud

0.14.0 e superior

Autenticação de ID do Google Cloud

0.14.0 e superior

Autenticação de access tokens pessoais do Databricks

Todas as versões

Etapa 1: Instalar ou atualizar o Databricks SDK for Python

  1. Os notebooks Databricks Python podem usar o Databricks SDK for Python como qualquer outra biblioteca do Python. Para instalar ou atualizar a biblioteca do Databricks SDK for Python no cluster do Databricks anexado, execute o comando mágico %pip em uma célula do notebook da seguinte forma:

    %pip install databricks-sdk --upgrade

  2. Depois de executar o comando mágico %pip, você deve reiniciar o Python para disponibilizar a biblioteca instalada ou atualizada para o notebook. Para fazer isto, execute o seguinte comando em uma célula do notebook imediatamente após a célula com o comando mágico %pip :

    dbutils.library.restartPython()

  3. Para exibir a versão instalada do Databricks SDK for Python, execute o seguinte comando em uma célula do notebook:

    %pip show databricks-sdk | grep -oP '(?<=Version: )\S+'

Etapa 2: Executar seu código

Nas células de seu notebook, crie o código do Python que importa e depois chama o Databricks SDK for Python. O exemplo a seguir usa a autenticação padrão do notebook do Databricks para listar todos os clusters em seu workspace do Databricks:

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

for c in w.clusters.list():
  print(c.cluster_name)

Quando você executa essa célula, aparece uma lista dos nomes de todos os clusters disponíveis no seu workspace do Databricks.

Para usar um tipo de autenticação diferente do Databricks, consulte Métodos de autenticação do Databricks e clique no link correspondente para obter detalhes técnicos adicionais.

Usar o Databricks Utilities

Você pode chamar a referência do Databricks Utilities (dbutils) do SDK do Databricks para código Python em execução em sua máquina de desenvolvimento local ou em um notebook Databricks.

  • Do seu computador de desenvolvimento local, o Databricks Utilities tem acesso somente aos grupos de comando dbutils.fs, dbutils.secrets, dbutils.widgets e dbutils.jobs.

  • Em um notebook do Databricks conectado a um cluster do Databricks, o Databricks Utilities tem acesso a todos os grupos de comandos disponíveis do Databricks Utilities, não apenas dbutils.fs, dbutils.secrets e dbutils.widgets. Além disso, o grupo de comandos dbutils.notebook está limitado a apenas dois níveis de comandos, por exemplo dbutils.notebook.run ou dbutils.notebook.exit.

Para chamar o Databricks Utilities de sua máquina de desenvolvimento local ou de um notebook do Databricks, use dbutils em WorkspaceClient. Este exemplo de código usa a autenticação padrão do notebook do Databricks para chamar dbutils em WorkspaceClient para listar os caminhos de todos os objetos na DBFS root do workspace.

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
d = w.dbutils.fs.ls('/')

for f in d:
  print(f.path)

Ou então, você pode chamar dbutils diretamente. No entanto, você está limitado a usar apenas a autenticação default do notebook do Databricks. Este exemplo de código chama dbutils diretamente para listar todos os objetos na DBFS root do workspace.

from databricks.sdk.runtime import *

d = dbutils.fs.ls('/')

for f in d:
  print(f.path)

Para acessar os volumes do Unity Catalog, use files dentro de WorkspaceClient. Consulte Gerenciar arquivos em volumes do Unity Catalog. Você não pode usar dbutils sozinho ou dentro de WorkspaceClient para acessar volumes.

Consulte também Interação com dbutils.

Exemplos de código

Os exemplos de código a seguir demonstram como usar o Databricks SDK for Python para criar e excluir clusters, executar jobs e listar grupos em nível de conta. Esses exemplos de código usam a autenticação default do notebook do Databricks. Para obter detalhes sobre a autenticação default do notebook do Databricks, consulte Usar o Databricks SDK for Python de um notebook do Databricks. Para obter detalhes sobre a autenticação default fora de notebooks, consulte Autenticar o Databricks SDK for Python com sua conta ou workspace do Databricks.

Para mais exemplos de código, consulte os exemplos no repositório Databricks SDK for Python no GitHub. Veja também:

Criar um cluster

Este exemplo de código cria um cluster com a versão especificada do Databricks Runtime e o tipo de nó de cluster. Este cluster tem um worker e será encerrado automaticamente após 15 minutos parado.

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

print("Attempting to create the cluster. Please wait...")

c = w.clusters.create_and_wait(
  cluster_name             = 'my-cluster',
  spark_version            = '12.2.x-scala2.12',
  node_type_id             = 'n2-highmem-4',
  autotermination_minutes = 15,
  num_workers              = 1
)

print(f"View the cluster at " \
      f"{w.config.host}#setting/clusters/{c.cluster_id}/configuration\n")

Excluir permanentemente um cluster

Este exemplo de código exclui permanentemente o cluster com o ID de cluster especificado do workspace.

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

c_id = input('ID of cluster to delete (for example, 1234-567890-ab123cd4): ')

w.clusters.permanent_delete(cluster_id = c_id)

Criar um job

Este exemplo de código cria um job do Databricks que executa o notebook especificado no cluster especificado. Conforme o código é executado, ele obtém o caminho do notebook existente, o ID do cluster existente e as configurações de job relacionadas do usuário no terminal.

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import Task, NotebookTask, Source

w = WorkspaceClient()

job_name            = input("Some short name for the job (for example, my-job): ")
description         = input("Some short description for the job (for example, My job): ")
existing_cluster_id = input("ID of the existing cluster in the workspace to run the job on (for example, 1234-567890-ab123cd4): ")
notebook_path       = input("Workspace path of the notebook to run (for example, /Users/someone@example.com/my-notebook): ")
task_key            = input("Some key to apply to the job's tasks (for example, my-key): ")

print("Attempting to create the job. Please wait...\n")

j = w.jobs.create(
  name = job_name,
  tasks = [
    Task(
      description = description,
      existing_cluster_id = existing_cluster_id,
      notebook_task = NotebookTask(
        base_parameters = dict(""),
        notebook_path = notebook_path,
        source = Source("WORKSPACE")
      ),
      task_key = task_key
    )
  ]
)

print(f"View the job at {w.config.host}/#job/{j.job_id}\n")

Gerenciar arquivos em volumes do Unity Catalog

Este exemplo de código demonstra várias chamadas da funcionalidade files em WorkspaceClient para acessar um volume do Unity Catalog.

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

# Define volume, folder, and file details.
catalog            = 'main'
schema             = 'default'
volume             = 'my-volume'
volume_path        = f"/Volumes/{catalog}/{schema}/{volume}" # /Volumes/main/default/my-volume
volume_folder      = 'my-folder'
volume_folder_path = f"{volume_path}/{volume_folder}" # /Volumes/main/default/my-volume/my-folder
volume_file        = 'data.csv'
volume_file_path   = f"{volume_folder_path}/{volume_file}" # /Volumes/main/default/my-volume/my-folder/data.csv
upload_file_path   = './data.csv'

# Create an empty folder in a volume.
w.files.create_directory(volume_folder_path)

# Upload a file to a volume.
with open(upload_file_path, 'rb') as file:
  file_bytes = file.read()
  binary_data = io.BytesIO(file_bytes)
  w.files.upload(volume_file_path, binary_data, overwrite = True)

# List the contents of a volume.
for item in w.files.list_directory_contents(volume_path):
  print(item.path)

# List the contents of a folder in a volume.
for item in w.files.list_directory_contents(volume_folder_path):
  print(item.path)

# Print the contents of a file in a volume.
resp = w.files.download(volume_file_path)
print(str(resp.contents.read(), encoding='utf-8'))

# Delete a file from a volume.
w.files.delete(volume_file_path)

# Delete a folder from a volume.
w.files.delete_directory(volume_folder_path)

Listar grupos em nível de conta

Este exemplo de código lista os nomes de exibição de todos os grupos disponíveis na conta do Databricks.

from databricks.sdk import AccountClient

a = AccountClient()

for g in a.groups.list():
  print(g.display_name)

Testando

Para testar seu código, use estruturas de teste Python como pytest. Para testar seu código em condições simuladas sem chamar endpoints da API REST do Databricks ou alterar o estado de suas accounts ou workspaces do Databricks, use bibliotecas de mocking do Python, como unittest.mock.

Dica

O Databricks Labs fornece um plug-in pytest para simplificar os testes de integração com o Databricks e um plug-in pylint para garantir a qualidade do código.

O arquivo de exemplo a seguir, denominado helpers.py, contém uma função create_cluster que retorna informações sobre o novo cluster:

# helpers.py

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.compute import ClusterDetails

def create_cluster(
  w: WorkspaceClient,
  cluster_name:            str,
  spark_version:           str,
  node_type_id:            str,
  autotermination_minutes: int,
  num_workers:             int
) -> ClusterDetails:
  response = w.clusters.create(
    cluster_name            = cluster_name,
    spark_version           = spark_version,
    node_type_id            = node_type_id,
    autotermination_minutes = autotermination_minutes,
    num_workers             = num_workers
  )
  return response

Dado o seguinte arquivo chamado main.py que chama a função create_cluster:

# main.py

from databricks.sdk import WorkspaceClient
from helpers import *

w = WorkspaceClient()

# Replace <spark-version> with the target Spark version string.
# Replace <node-type-id> with the target node type string.
response = create_cluster(
  w = w,
  cluster_name            = 'Test Cluster',
  spark_version           = '<spark-version>',
  node_type_id            = '<node-type-id>',
  autotermination_minutes = 15,
  num_workers             = 1
)

print(response.cluster_id)

O seguinte arquivo chamado test_helpers.py testa se a função create_cluster retorna a resposta esperada. Em vez de criar um cluster no workspace de destino, este teste simula um objeto WorkspaceClient, define as configurações do objeto simulado e, em seguida, passa o objeto simulado para a função create_cluster. O teste verifica se a função retorna o ID esperado do novo cluster simulado.

# test_helpers.py

from databricks.sdk import WorkspaceClient
from helpers import *
from unittest.mock import create_autospec # Included with the Python standard library.

def test_create_cluster():
  # Create a mock WorkspaceClient.
  mock_workspace_client = create_autospec(WorkspaceClient)

  # Set the mock WorkspaceClient's clusters.create().cluster_id value.
  mock_workspace_client.clusters.create.return_value.cluster_id = '123abc'

  # Call the actual function but with the mock WorkspaceClient.
  # Replace <spark-version> with the target Spark version string.
  # Replace <node-type-id> with the target node type string.
  response = create_cluster(
    w = mock_workspace_client,
    cluster_name            = 'Test Cluster',
    spark_version           = '<spark-version>',
    node_type_id            = '<node-type-id>',
    autotermination_minutes = 15,
    num_workers             = 1
  )

  # Assert that the function returned the mocked cluster ID.
  assert response.cluster_id == '123abc'

Para executar este teste, execute o comando pytest a partir da raiz do projeto de código, o que deve produzir resultados de teste semelhantes aos seguintes:

$ pytest
=================== test session starts ====================
platform darwin -- Python 3.12.2, pytest-8.1.1, pluggy-1.4.0
rootdir: <project-rootdir>
collected 1 item

test_helpers.py . [100%]
======================== 1 passed ==========================