Unity Catalog テーブルを PyTorch にロードする UCDataset
ベータ版
この機能はベータ版です。
UCDataset は、AI ランタイムのトレーニングループにUnity CatalogテーブルからデータをストリームするPyTorch IterableDataset です。DataFrameとして読み込む対象（テーブルの選択、列の選択、行のフィルタリングなど）を記述します。データセットは、クラウドストレージのテーブルのデータファイルから一致する行を直接読み込むため、SparkセッションやSQL Warehouseは不要です。読み込み時にテーブルの削除ベクトルを適用し、一致する行をGPUとデータローダーワーカー全体に均等に分割し、バッチごとに1つの辞書を生成します。数値列とBoolean列はテンソルになります。他の列はPythonリストになります。
トレーニングデータがどこにあるかに基づいてデータセットを選択します。
学習用データ | データセット |
|---|---|
Delta Lakeテーブルの特徴量列やラベル列など、Unity Catalogテーブルの行 |
|
画像、音声、ドキュメントなど、Unity Catalog ボリューム内のファイル |
2つのデータセットは並行して機能します。たとえば、トレーニングジョブは、UCDatasetを含むテーブルからラベルを読み取り、UCVolumeDatasetを含むボリュームから対応する画像を読み取ることができます。
その他のデータについては、代わりにこれらのオプションを使用してください。
- メモリに収まる小さなテーブルや、結合や集計が必要なクエリーの場合は、Spark Connectを使用してください。
- Ray ワークロードについては、「
ray.data.read_databricks_tables()」を参照してください。
要件
- AI ランタイムワークロード。
UCDatasetは、AI ランタイムがある場所であればどこでも利用できます。「 Geo 別の AI ランタイムの可用性」を参照してください。 - Python 3.10 以降がインストールされている x86_64 版 Linux 。このパッケージは x86_64 専用です。AI ランタイムの GPU コンピュートはこの要件を満たしています。
- Unity Catalogに登録されているDelta Lakeテーブル。
- テーブルのカタログに対する
USE CATALOG特権、そのスキーマに対するUSE SCHEMA特権、およびテーブルに対するSELECT特権。 - 2つ以上のGPUでトレーニングするには、
torchrunのように、すべてのプロセスに対してRANK環境変数とWORLD_SIZE環境変数を設定するランチャーが必要です。AI ランタイムノートブックのように変数がどちらも設定されていない場合、UCDatasetはテーブルを単一のシャードとして読み取ります。変数のうち1つだけが設定されている場合、UCDatasetはエラーを発生させます。
databricks air run でワークロードをサブミットする場合、シングルGPUであっても torchrun でトレーニング スクリプトを起動してください。AIランタイムはワークロードごとに WORLD_SIZE を設定しますが、 RANK は設定しません。 UCDataset は両方のいずれかのみが設定されている場合にエラーをスローします。torchrun は各プロセスの両方を設定します。AI ランタイムでのマルチノード ジョブのランを参照してください。
テーブルの読み込み
-
databricks-uc-datasetパッケージをインストールします。databricks air runで送信するワークロードの場合、パッケージをenvironment.dependenciesに追加します。「 AI ランタイムでマルチノードジョブをランする」を参照してください。ノートブックで、以下をランします。Python%pip install databricks-uc-dataset==0.0.1b1 -
読み取るテーブル、列、および行を選択します。
table()は3レベルのテーブル名を受け取り、オプションのversion_as_of引数は特定のテーブルバージョンを読み取ります：Pythonfrom databricks.uc_dataset.beta.dataframe import table, coldf = (table("main.ml.clicks").select("features", "label").filter(col("label").isNotNull() & col("features").isNotNull())) -
データセットを作成し、PyTorch
DataLoaderでラップします。データセットはすでにバッチを返しているため、ローダーでbatch_size=Noneを設定します。PyTorch DistributedDataParallel (DDP) でトレーニングする場合は、データセットを作成する前にtorch.distributed.init_process_group()を呼び出します。ランク 0 は読み取りを 1 回計画し、その計画を他のランクと共有します。Pythonimport torchfrom torch.utils.data import DataLoaderfrom databricks.uc_dataset.beta import UCDatasetdataset = 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) -
1 つ以上の GPU でトレーニングするには、コードを
train.pyとして保存し、torchrunで実行します。ノートブックでは、このステップをスキップします。たとえば、1 つのノードで 8 つの GPU でトレーニングするには、以下を実行します。Bashtorchrun --nproc_per_node=8 train.py
各ランクとワーカーは、一致する行の独自のシャードを読み取ります。各バッチは、選択された各列名をTensorまたはリストにマップする辞書です。
AIランタイムでマルチノード ジョブをランします
複数のノードでトレーニングする場合は、ワークロードを workload.yaml で記述し、 databricks air run でサブミットします。
-
environment.versionをdatabricks_ai_v6に設定します（これには PyTorch が含まれます）。また、databricks-uc-dataset==0.0.1b1をenvironment.dependenciesに追加します。 -
torchrunでトレーニングスクリプトを起動します。AI ランタイムはすべてのノードでコマンドを実行し、NUM_NODES、NODE_RANK、LOCAL_WORLD_SIZE、MASTER_ADDR、およびMASTER_PORTを設定します。torchrunはこれらを使用して、プロセスごとにRANKを設定します。YAMLcommand: |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"
train.py、workload.yaml、およびサンプルテーブルを使用した完全なワークロードについては、AI ランタイム Cookbookの「 Unity Catalog テーブルのクリックスルー率モデル」を参照してください。より多くの GPU でトレーニングするには、workload.yaml の num_accelerators を変更します。train.py は変更されません。
行のフィルター
UCDataset はファイルを読み取るときにフィルターを適用するため、一致しない行はトレーニングループに到達しません。col() および以下の式を使用してフィルターを構築します。
式 | 例 |
|---|---|
比較 |
|
Null値のチェック |
|
Nullセーフな等価性 |
|
Boolean論理 |
|
フィルターでは、SQL文字列、isinまたはIN、BETWEEN、likeやcontainsなどの文字列照合、列に対する算術演算、cast、ユーザー定義関数（UDF）はサポートされません。各条件をColumn式として渡し、&、|、~で条件を組み合わせます。比較では、列とリテラル値をどちらの順序でも指定できます。2つの列を比較することはできません。値のリストに一致させるには、==条件を|で組み合わせます。フィルターが非常に長いと、失敗することがあります。制限事項を参照してください。
次の例は、固定テーブルバージョンから購入イベントを読み取ります。バージョンが存在する必要があります：
df = (
table("main.ml.events", version_as_of=120)
.filter((col("event_type") == "purchase") & (col("amount") > 50))
.select("user_id", "amount", "label")
)
出力フォーマット
各バッチは、選択した各列名をその値にマッピングするPython辞書です。
列タイプ | バッチ値 |
|---|---|
整数（バッチにnullなし） | 列のInteger型の1次元テンソル（以下など） |
整数。バッチ内にNULLがある場合（64ビットを除く） | 1次元 |
浮動小数点 | 1次元テンソル。NULLは次として扱われます |
Boolean（Nullはなし） | Booleanテンソル |
Binary |
|
文字列、配列、構造体、Timestampなどのその他の型 | Pythonリスト |
整数列の型はバッチに null が含まれるかどうかによって異なるため、列の型はバッチごとに変わる可能性があります。モデルが埋め込み検索やクラスラベルなどの固定Integer型を必要とする場合は、null を除外するか、transform 関数で置き換えます。
データを失わずに変換できない列の場合、UCDatasetはエラーをスローします。これは、NULLを含む64ビット整数列と、NULLを含むBoolean列で発生します。エラーはNULLを含むバッチが読み取られたときに発生するため、エポックの途中で発生する可能性があります。これを避けるには、isNotNull()でNULLをフィルタリングするか、transform関数でNULLを置き換えます。
バッチのLayoutを変更する場合、またはNULLを処理する場合は、transform関数を渡します。UCDatasetは、テンソルに変換する前に、バッチ内の各要素をPyArrow RecordBatchとして呼び出し、その関数が返すものをすべて生成します。使用するすべての列を変換してください。次の例では、テンソルのタプルを返します。
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)
次の例では、64ビット整数列のnullを0に置き換えるため、列はすべてのバッチでint64のままです。
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)
分散トレーニングとシャッフル
UCDataset フィルターに一致する行を、すべてのランクとすべてのデータローダーワーカーに分割します。分散サンプラーは必要ありません。defaultでは、シャードサイズは最大で1行異なります。
シャードサイズが等しいからといって、バッチ数が等しいとは限りません。あるランクは別のランクより1つ多くのバッチを取得できる場合があります。例えば、いくつかのシャードがbatch_size行の正確な倍数を保持し、他が1行多く保持している場合や、フィルターがシャードよりも少ない行を保持する場合などです。DDPでは、追加のバッチを持つランクは、他のランクが決して参加しない勾配のオールリデュースを待機し、ジョブは失敗します。ランクごとにエポックを異なるタイミングで終了させるには、各エポックをPyTorchのJoinコンテキストマネージャーでラップします。ここでmodelはDDPでラップされたモデルです：
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() フィルターに一致する行数を、すべてのランクとワーカーにわたって返します。len(dataset)とlen(loader)はTypeErrorをスローします。ライブラリに学習率スケジューラのようなエポックのステップ数が必要な場合は、num_rows()、ワールドサイズ、およびbatch_sizeから推定します。各ランクは、推定値より1つ多くのバッチを読み取ることができます。また、各追加データローダーワーカーは、別の部分バッチを追加できます。
UCDataset テーブルのファイルを固定されたシャッフルされた順序で読み取ります。各エポックでは同じ順序が使用され、変更されていないテーブル上の新しいデータセットでも同様です。シードを設定できず、エポック間でファイルが再シャッフルされることもありません。
認証
UCDataset ワークロードをランする ID としてテーブルを読み込み、テーブルに対する Unity Catalog の権限が適用されます。
Unity Catalog は、行フィルターまたは列マスクを持つテーブルのストレージ資格情報を発行しないため、UCDataset はそれらを読み取ることができません。データセットの作成は、Unity Catalog からのメッセージを含むエラーで失敗します。
制限
UCDatasetUnity Catalog の Delta Lake テーブルのみを読み込みます。ビュー、マテリアライズドビュー、またはストリーミングテーブルはサポートされていません。- 行フィルターまたは列マスクを持つテーブルは読み取ることができません。これは Unity Catalog がそれらのストレージ資格情報を発行しないためです。
- 浅いクローンは読み取れません。データファイルがソーステーブルとともに保存されているためです。トレーニングランの行をピン留めするには、代わりに
version_as_ofを使用してソーステーブルを読み取ります。クローンは独自のバージョン履歴を起動するため、ソーステーブルの以前のバージョンは含まれません。 VARIANT列を含むテーブルは、その列を選択しない場合でも読み取れません。version_as_ofテーブルのバージョン番号を指定します。Timestampはサポートされていません。version_as_of="2026-01-01"はversionAsOf must be a non-negative integerエラーを発生させます。- パーティション列ではない列を少なくとも 1 つ選択してください。パーティション列のみを選択すると、
Delta reads require at least one non-partition columnエラーが発生します。 - 多くの
==の比較を|で結合したような、条件の長いチェーンを持つフィルターは、failed to decode Spark Connect planエラーで失敗する可能性があります。その場合は、条件の数を減らしてください。 - フィルターは、フィルター行の式のみをサポートしています。
- 埋め込みなどの配列カラムは Python のリストとして返されるため、数値カラムよりも低速です。
UCDataset読み取り位置をチェックポイントとして保存しません。トレーニングが再開されると、データセットはテーブルの先頭から再び開始され、エポックの途中から再開することはできません。
ベータ期間中に問題を報告またはフィードバックを提供するには、Databricksアカウントチームに連絡してください。