Carregar tabelas do Unity Catalog no PyTorch com `UCDataset`

info

Beta

Este recurso está em Beta.

UCDataset é um PyTorch IterableDataset que transmite dados de tabelas do Unity Catalog para um loop de treinamento no AI Runtime. Você descreve o que ler como um DataFrame: escolha uma tabela, selecione as colunas e filtre as linhas. O dataset lê as linhas correspondentes diretamente dos arquivos de dados da tabela no cloud storage, portanto, você não precisa de uma sessão Spark ou de um SQL warehouse. Ele aplica os vetores de exclusão da tabela durante a leitura, divide as linhas correspondentes uniformemente entre GPUs e worker do carregador de dados e produz um dicionário por lotes. As colunas numéricas e Boolean tornam-se tensores. Outras colunas tornam-se listas Python.

Escolha um dataset com base em onde seus dados de treinamento estão localizados:

Dados de treinamento	Conjunto de dados
Linhas em uma tabela do Unity Catalog, como colunas de recurso e rótulo em uma tabela Delta Lake	`UCDataset` (Beta)
Arquivos em um volume do Unity Catalog, como imagens, áudio ou documentos	`UCVolumeDataset`

Dados de treinamento	Conjunto de dados
Linhas em uma tabela do Unity Catalog, como colunas de recurso e rótulo em uma tabela Delta Lake	`UCDataset` (Beta)
Arquivos em um volume do Unity Catalog, como imagens, áudio ou documentos	`UCVolumeDataset`

Os dois datasets funcionam lado a lado. Por exemplo, um job de treinamento pode ler rótulos de uma tabela com UCDataset e as imagens correspondentes de um volume com UCVolumeDataset.

Para outros dados, use estas opções:

Para tabelas pequenas que cabem na memória, ou para query que precisam de joins ou agregações, use o Spark Connect.
Para cargas de trabalho Ray, use ray.data.read_databricks_tables().

Requisitos

Uma carga de trabalho de Runtime de AI. UCDataset está disponível onde quer que o AI Runtime esteja. Consulte Disponibilidade do Runtime de AI por Geo.
Linux em x86_64 com Python 3.10 ou acima. O pacote foi compilado apenas para x86_64. O compute GPU de Runtime de AI atende a esse requisito.
Uma tabela Delta Lake registrada no Unity Catalog.
O privilégio USE CATALOG no catálogo da tabela, o privilégio USE SCHEMA em seu esquema e o privilégio SELECT na tabela.
Para ensinar em mais de uma GPU, um iniciador que define as variáveis de ambiente RANK e WORLD_SIZE para cada processo, como torchrun. Se nenhuma variável estiver definida, como em um notebook do AI Runtime, UCDataset lê a tabela como um único shard. Se apenas um deles estiver definido, UCDataset levantará um erro.

importante

Quando você envia uma carga de trabalho com databricks air run, começar o script de treinamento com torchrun, mesmo em uma única GPU. O AI Runtime define WORLD_SIZE para cada carga de trabalho, mas não define RANK, e UCDataset gera um erro quando apenas um dos dois é definido. torchrun define ambos para cada processo. Consulte execução um Job multi-nó no AI Runtime.

Carregar uma tabela

Instale o pacote databricks-uc-dataset . Para uma carga de trabalho que você envia com databricks air run, adicione o pacote a environment.dependencies. Consulte execução um Job multi-nó no AI Runtime. Em um notebook, execução o seguinte:
Python
```
%pip install databricks-uc-dataset==0.0.1b1
```

Selecione a tabela, colunas e linhas para ler. table() aceita um nome de tabela de três níveis, e o argumento opcional version_as_of lê uma versão de tabela específica:

Python
from databricks.uc_dataset.beta.dataframe import table, col

df = (
    table("main.ml.clicks")
    .select("features", "label")
    .filter(col("label").isNotNull() & col("features").isNotNull())
)

Crie o dataset e envolva-o em um PyTorch DataLoader. O dataset já retorna lotes, então defina batch_size=None no loader. Se você treinar com o PyTorch DistributedDataParallel (DDP), chame torch.distributed.init_process_group() antes de criar o dataset. O ranking 0 então planeja a leitura uma vez e compartilha o plano com os outros rankings:

Python
import torch
from torch.utils.data import DataLoader
from databricks.uc_dataset.beta import UCDataset

dataset = UCDataset(df=df, batch_size=256)
loader = DataLoader(dataset, batch_size=None, num_workers=4)

for epoch in range(num_epochs):
    for batch in loader:
        # Array columns arrive as Python lists, and a NULL array arrives as None.
        # Convert the lists before you train.
        features = torch.as_tensor(batch["features"], dtype=torch.float32)
        label = batch["label"]
        train_step(features, label)

Para ensinar em mais de uma GPU, salve o código como train.py e inicie-o com torchrun. Em um notebook, ignore o passo. Por exemplo, para ensinar em 8 GPUs em um nó, execute o seguinte:
Bash
```
torchrun --nproc_per_node=8 train.py
```

Cada rank e worker lê seu próprio shard das linhas correspondentes. Cada lote é um dicionário que mapeia cada nome de coluna selecionado para um tensor ou uma lista.

Execução de um Job multi-nó no AI Runtime

Para ensinar em mais de um nó, descreva a carga de trabalho em workload.yaml e envie-a com databricks air run:

Defina environment.version como databricks_ai_v6, que inclui o PyTorch, e adicione databricks-uc-dataset==0.0.1b1 a environment.dependencies.
Comece o script de treinamento com torchrun. O AI Runtime executa o comando em cada nó e define NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR e MASTER_PORT. torchrun os usa para definir RANK para cada processo:
YAML
```
command: |
  torchrun --nnodes=$NUM_NODES --node_rank=$NODE_RANK --nproc_per_node=$LOCAL_WORLD_SIZE \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT "$CODE_SOURCE_PATH/train.py"
```

Para uma carga de trabalho completa com train.py, workload.yaml e uma tabela de amostra, consulte Modelo de taxa de cliques em uma tabela do Unity Catalog no Livro de Receitas do AI Runtime. Para ensinar em mais GPUs, altere num_accelerators em workload.yaml. train.py não muda.

Filtrar linhas

UCDataset aplica filtros enquanto lê os arquivos, então as linhas que não correspondem não alcançam seu loop de treinamento. Crie um filtro com col() e as seguintes expressões:

Expressão	Exemplo
Comparação	`col("amount") > 50`. Também `==`, `!=`, `<`, `<=` e `>=`.
Verificação de nulo	`col("label").isNull()`, `col("label").isNotNull()`
Igualdade segura para nulos	`col("region").eqNullSafe("us")`
Lógica Boolean	`(col("a") > 1) & (col("b") == "x")`. Além disso `\|` (ou) e `~` (não).

Expressão	Exemplo
Comparação	`col("amount") > 50`. Também `==`, `!=`, `<`, `<=` e `>=`.
Verificação de nulo	`col("label").isNull()`, `col("label").isNotNull()`
Igualdade segura para nulos	`col("region").eqNullSafe("us")`
Lógica Boolean	`(col("a") > 1) & (col("b") == "x")`. Além disso `\|` (ou) e `~` (não).

Os filtros não suportam strings SQL, isin ou IN, BETWEEN, correspondência de strings como like ou contains, aritmética em colunas, cast ou funções definidas pelo usuário (UDFs). Passe cada condição como uma expressão Column e combine condições com &, | e ~. Uma comparação pega uma coluna e um valor literal, em qualquer ordem. Não é possível comparar duas colunas. Para corresponder a uma lista de valores, combine as condições == com |. Filtros muito longos podem falhar. Consulte Limitações.

O exemplo a seguir lê eventos de compra de uma versão fixa de tabela. A versão deve existir:

Python
df = (
    table("main.ml.events", version_as_of=120)
    .filter((col("event_type") == "purchase") & (col("amount") > 50))
    .select("user_id", "amount", "label")
)

Formato de saída

Cada lote é um dicionário Python que mapeia cada nome de coluna selecionado para seus valores:

Tipo da coluna	Valor de lote
Inteiro, sem nulos no lote	Tensor 1-D com o tipo de integer types da coluna, como `int32`
Inteiro, com nulos no lote (exceto 64 bits)	Tensor `float64` unidimensional, com nulos como `NaN`
Ponto flutuante	Tensor 1-D, com nulos como `NaN`
Boolean, sem nulos	Tensor Boolean
binário	Lista de `bytes` objetos
Outros tipos, como strings, arrays, structs e Timestamp	Lista Python

Tipo da coluna	Valor de lote
Inteiro, sem nulos no lote	Tensor 1-D com o tipo de integer types da coluna, como `int32`
Inteiro, com nulos no lote (exceto 64 bits)	Tensor `float64` unidimensional, com nulos como `NaN`
Ponto flutuante	Tensor 1-D, com nulos como `NaN`
Boolean, sem nulos	Tensor Boolean
binário	Lista de `bytes` objetos
Outros tipos, como strings, arrays, structs e Timestamp	Lista Python

Como o tipo de uma coluna inteira depende se o lote contém um nulo, a coluna pode mudar o tipo de um lote para o próximo. Se seu modelo precisar de integer types fixo, como para consultas de incorporação ou rótulos de classe, filtre os nulos ou substitua-os em uma função transform.

Se uma coluna não puder ser convertida sem perder dados, UCDataset gera um erro. Isso acontece para uma coluna de inteiro de 64 bits que contém nulos e para uma coluna Boolean que contém nulos. O erro aparece quando um lote com um nulo é lido, então pode acontecer no meio de uma época. Para evitar isso, filtre os nulos com isNotNull() ou substitua-os em uma função transform.

Para alterar a disposição do lote ou lidar com nulos, passe uma função transform. UCDataset o chama em cada lote como um RecordBatch PyArrow, antes de qualquer conversão para tensores, e gera o que a função retornar, portanto converta todas as colunas que você usar. O exemplo a seguir retorna uma tupla de tensores:

Python
import torch

def to_features_and_label(batch):
    features = torch.as_tensor(batch.column("features").to_pylist(), dtype=torch.float32)
    label = torch.from_numpy(batch.column("label").to_numpy(zero_copy_only=False))
    return features, label

dataset = UCDataset(df=df, transform=to_features_and_label)

O exemplo a seguir substitui nulos em uma coluna de inteiro de 64 bits por 0, portanto, a coluna permanece int64 em cada lote:

Python
import pyarrow.compute as pc
import torch

def fill_nulls(batch):
    user_ids = pc.fill_null(batch.column("user_id"), 0)
    return {"user_id": torch.tensor(user_ids.to_numpy())}

dataset = UCDataset(df=df, transform=fill_nulls)

Treinamento e embaralhamento distribuídos

UCDataset divide as linhas que correspondem ao seu filtro entre cada rank e cada worker do carregador de dados. Você não precisa de um amostrador distribuído. Por default, os tamanhos de shard diferem em no máximo uma linha.

Tamanhos de shard iguais não garantem contagens de lote iguais. Uma classificação pode obter um lote a mais do que outra, por exemplo, quando alguns shards contêm um múltiplo exato de batch_size linhas e outros contêm uma linha a mais, ou quando um filtro mantém menos linhas do que existem shards. Com DDP, a classificação com os lotes extra aguarda uma redução total de gradiente que as outras classificações nunca join, e o job falha. Para permitir que as classificações finalizem uma época em momentos diferentes, envolva cada época no gerenciador de contexto Join do PyTorch, onde model é o modelo envolto em DDP:

Python
from torch.distributed.algorithms.join import Join

for epoch in range(num_epochs):
    with Join([model]):
        for batch in loader:
            train_step(batch)

dataset.num_rows() retorna o número de linhas que correspondem ao filtro, em todos os ranks e workers. len(dataset) e len(loader) geram um TypeError. Se uma biblioteca precisar do número de os passos em uma época, como um programador de taxa de aprendizado, estime-o a partir de num_rows(), o tamanho do mundo e batch_size. Cada rank pode ler mais um lote do que a estimativa, e cada worker de carregador de dados adicional pode adicionar outro lote parcial.

UCDataset lê os arquivos da tabela em uma ordem fixa e embaralhada. Cada época usa a mesma ordem, assim como um novo dataset em uma tabela não alterada. Você não pode definir uma seed, e os arquivos não são reembaralhados entre as épocas.

Autenticação

UCDataset lê a tabela como a identidade que executa a carga de trabalho, e as permissões do Unity Catalog na tabela são aplicadas.

O Unity Catalog não emite credenciais de armazenamento para tabelas com filtros de linha ou máscaras de coluna, portanto UCDataset não pode lê-las. A criação do dataset falha com um erro que inclui a mensagem do Unity Catalog.

Limitações

UCDataset lê apenas tabelas Delta Lake no Unity Catalog. Views, materialized views e tabelas de transmissão não são suportadas.
Tabelas com filtros de linha ou máscaras de coluna não podem ser lidas, porque o Unity Catalog não emite credenciais de armazenamento para elas.
Clones superficiais não podem ser lidos, porque seus arquivos de dados são armazenados com a tabela de origem. Para pin as linhas para uma execução de treinamento, leia a tabela de origem com version_as_of em vez disso. Um clone inicia sua própria história de versões, então ele não possui as versões anteriores da tabela de origem.
Tabelas com uma coluna VARIANT não podem ser lidas, mesmo que você não selecione essa coluna.
version_as_of pega um número de versão da tabela. Timestamp não são suportados: version_as_of="2026-01-01" gera um erro versionAsOf must be a non-negative integer.
Selecione pelo menos uma coluna que não seja uma coluna de partição. Selecionar apenas colunas de partição gera um erro Delta reads require at least one non-partition column.
Um filtro com uma longa cadeia de condições, como muitas comparações == combinadas com |, pode falhar com um erro failed to decode Spark Connect plan. Se acontecer, use menos condições.
Os filtros suportam apenas as expressões em Filtrar linhas.
As colunas de matriz, como embeddings, são retornadas como listas Python, o que é muito mais lento do que as colunas numéricas.
UCDataset não verifica sua posição de leitura. Se o treinamento reiniciar, o dataset começa novamente do início da tabela, e você não pode retomar no meio de uma época.

Para relatar um problema ou dar feedback durante a Beta, entre em contato com sua equipe de account Databricks.

Requisitos​

Carregar uma tabela​

Execução de um Job multi-nó no AI Runtime​

Filtrar linhas​

Formato de saída​

Treinamento e embaralhamento distribuídos​

Autenticação​

Limitações​

Outros recursos​