kaxil commented on code in PR #76317: URL: https://github.com/apache/airflow/pull/76317#discussion_r4524210322
########## providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py: ########## @@ -0,0 +1,796 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Operators to load, transform, and publish Hugging Face datasets.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import marshal +import tempfile +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from datasets import Dataset, DatasetDict + +from airflow.providers.common.compat.sdk import Asset, BaseOperator, Context, ObjectStoragePath +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook + + +def _path(uri: str, *, conn_id: str | None = None) -> ObjectStoragePath: + """Resolve a URI through Airflow's configured object-storage filesystems.""" + return ObjectStoragePath(uri, conn_id=conn_id) + + +def _reference(reference: str | dict[str, Any]) -> str: + if isinstance(reference, dict): + uri = reference.get("uri") + if not uri: + raise ValueError("Artifact reference must contain a non-empty 'uri'.") + return uri + if not isinstance(reference, str) or not reference: + raise ValueError("Artifact reference must be a URI or contain a 'uri'.") + return reference + + +class _HashingWriter: + def __init__(self, stream: Any) -> None: + self._stream = stream + self._digest = hashlib.sha256() + + def write(self, data: bytes) -> int: + written = self._stream.write(data) + self._digest.update(data[:written]) + return written + + def get_hexdigest(self) -> str: + return self._digest.hexdigest() + + def __getattr__(self, name: str) -> Any: + return getattr(self._stream, name) + + +def _validate_feature(feature: Any, field_name: str) -> None: + """Allow Parquet-safe scalar, class-label, and nested sequence features.""" + from datasets import ClassLabel, LargeList, Sequence, Value + + if isinstance(feature, (Value, ClassLabel)): + return + if isinstance(feature, (Sequence, LargeList)): + _validate_feature(feature.feature, field_name) + return + if isinstance(feature, dict): + for key, value in feature.items(): + _validate_feature(value, f"{field_name}.{key}") + return + if isinstance(feature, (list, tuple)): + for value in feature: + _validate_feature(value, field_name) + return + raise ValueError( + f"unsupported Parquet feature type for {field_name!r}: {type(feature).__name__}; " + "supported types are Value, ClassLabel, and nested Sequence or LargeList features" + ) + + +def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str, Dataset]]: + from datasets import Dataset, DatasetDict + + if isinstance(dataset, DatasetDict): + return list(dataset.items()) + if isinstance(dataset, Dataset): + return [("dataset", dataset)] + raise TypeError("Expected a datasets.Dataset or datasets.DatasetDict.") + + +def _save( + dataset: Dataset | DatasetDict, + uri: str, + metadata: dict[str, Any] | None = None, + *, + storage_conn_id: str | None = None, +) -> dict[str, Any]: + """Write split Parquet objects and a JSON manifest through ObjectStoragePath.""" + import pyarrow.parquet as parquet + + parts = _dataset_parts(dataset) + written: list[ObjectStoragePath] = [] + manifest_path = _path(uri, conn_id=storage_conn_id) + split_manifest: dict[str, Any] = {} + try: + for index, (name, part) in enumerate(parts): + for field_name, feature in part.features.items(): + _validate_feature(feature, field_name) + split_uri = f"{uri}.split-{index:05d}.parquet" + split_path = _path(split_uri, conn_id=storage_conn_id) + split_path.parent.mkdir(parents=True, exist_ok=True) + with split_path.open("wb") as stream: + hashing_stream = _HashingWriter(stream) + parquet.write_table(part.data.table, hashing_stream) + written.append(split_path) + split_manifest[name] = { + "uri": split_uri, + "rows": len(part), + "features": part.features.to_dict(), + "feature_order": list(part.features), + "split": ( + str(part.split) + if getattr(part, "split", None) is not None + else (name if name != "dataset" else None) + ), + "sha256": hashing_stream.get_hexdigest(), + } + + manifest = { + "format": "parquet", + "version": 1, + "splits": split_manifest, + "metadata": metadata or {}, + } + manifest_bytes = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest_path.parent.mkdir(parents=True, exist_ok=True) + with manifest_path.open("wb") as stream: + stream.write(manifest_bytes) + return { + "uri": uri, + "format": "parquet", + "splits": { + name: { + "rows": part_info["rows"], + "features": part_info["features"], + "split": part_info["split"], + } + for name, part_info in split_manifest.items() + }, + **(metadata or {}), + "fingerprint": hashlib.sha256(manifest_bytes).hexdigest(), + } + except Exception: + # Remove partial split outputs. A failed write must not leave an artifact + # reference that a downstream task could mistake for a complete dataset. + for path in written: + with suppress(OSError, NotImplementedError): + path.unlink(missing_ok=True) + with suppress(OSError, NotImplementedError): + manifest_path.unlink(missing_ok=True) + raise + + +def _read_artifact_manifest( + reference: str | dict[str, Any], *, storage_conn_id: str | None = None +) -> tuple[str, dict[str, Any]]: + uri = _reference(reference) + manifest_path = _path(uri, conn_id=storage_conn_id) + with manifest_path.open("rb") as stream: + manifest_bytes = stream.read() + if isinstance(reference, dict) and reference.get("fingerprint"): + actual = hashlib.sha256(manifest_bytes).hexdigest() + if actual != reference["fingerprint"]: + raise ValueError("Dataset artifact fingerprint does not match its reference") + + manifest = json.loads(manifest_bytes) + if manifest.get("format") != "parquet" or manifest.get("version") != 1: + raise ValueError(f"Unsupported dataset artifact manifest: {uri}") + return uri, manifest + + +def _load_artifact( + reference: str | dict[str, Any], + *, + cache_dir: str | None = None, + manifest: dict[str, Any] | None = None, + storage_conn_id: str | None = None, +) -> Dataset | DatasetDict: + from datasets import Dataset, DatasetDict, Features + + uri, manifest = ( + _read_artifact_manifest(reference, storage_conn_id=storage_conn_id) + if manifest is None + else (_reference(reference), manifest) + ) + resolved_cache_dir = cache_dir or str(Path(tempfile.gettempdir()) / "airflow-hf-datasets-cache") + Path(resolved_cache_dir).mkdir(parents=True, exist_ok=True) + parts: dict[str, Dataset] = {} + with tempfile.TemporaryDirectory(prefix="airflow-hf-parquet-") as local_dir: + for index, (name, item) in enumerate(manifest["splits"].items()): + split_uri = item["uri"] + local_path = Path(local_dir) / f"split-{index:05d}.parquet" + digest = hashlib.sha256() + with ( + _path(split_uri, conn_id=storage_conn_id).open("rb") as source, + local_path.open("wb") as target, + ): + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + target.write(chunk) + if digest.hexdigest() != item.get("sha256"): + raise ValueError(f"Dataset artifact split fingerprint mismatch: {split_uri}") + + feature_data = item["features"] + feature_order = item.get("feature_order", list(feature_data)) + features = Features.from_dict({key: feature_data[key] for key in feature_order}) + if item.get("rows") == 0: + restored = Dataset.from_dict({key: [] for key in features}, features=features) + else: + restored = Dataset.from_parquet( Review Comment: Without `split=`, `from_parquet` labels the restored Dataset `train` (checked on datasets 5.1.0), so the manifest's `split` is dropped here. A Load with `split="validation"` followed by a Push with no `split` in `push_kwargs` publishes those rows as `train`, since `push_to_hub` falls back to `self.split`. `test_single_dataset_split_name_survives_artifact_round_trip` fails for the same reason. Could the stored split be passed back when it's a plain name (`^\w+(\.\w+)*$`), and stay metadata only for expressions like `train+test`? ########## providers/huggingface/tests/unit/huggingface/operators/test_datasets.py: ########## @@ -0,0 +1,712 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Regression coverage for durable load/map/publish behavior.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from datasets import ClassLabel, Dataset, DatasetDict, Features, Image, LargeList, Value +from huggingface_hub import CommitInfo, GitCommitInfo + +from airflow.providers.huggingface.operators.datasets import ( + LoadDatasetOperator, + MapDatasetOperator, + PushToHubOperator, + StreamDatasetOperator, + _is_hub_repo, + _load_artifact, + _publication_marker, + _save, +) + + +def add_one(example): + return {"value": example["value"] + 1} + + +def commit_info(repo_id="org/data", oid="deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", message="publish"): + url = f"https://huggingface.co/datasets/{repo_id}/commit/{oid}"; + return CommitInfo(commit_url=url, commit_message=message, commit_description="", oid=oid) + + +def git_commit_info(commit_id, title, message=""): + return GitCommitInfo( + commit_id=commit_id, + authors=["unit-test"], + created_at=datetime(2026, 1, 1, tzinfo=UTC), + title=title, + message=message, + formatted_title=None, + formatted_message=None, + ) + + [email protected] +def hub_dependencies(monkeypatch): + """Use a real Airflow Connection while mocking only Hub and Datasets APIs.""" + monkeypatch.setenv( + "AIRFLOW_CONN_HUGGINGFACE_TEST", + '{"conn_type":"huggingface","password":"operator-test-token"}', + ) + with ( + patch("airflow.providers.huggingface.hooks.huggingface.HfApi", autospec=True) as hf_api, + patch("datasets.Dataset.push_to_hub", autospec=True) as push_to_hub, + ): + hf_api.return_value.list_repo_commits.return_value = [] + yield hf_api, push_to_hub + + +def test_load_pins_and_records_resolved_revision_and_honors_download_mode(tmp_path, monkeypatch): + output = (tmp_path / "loaded.json").as_uri() + dataset = Dataset.from_dict({"value": [1, 2]}) + monkeypatch.setenv( + "AIRFLOW_CONN_HUGGINGFACE_TEST", '{"conn_type":"huggingface","password":"operator-test-token"}' + ) + with ( + patch("airflow.providers.huggingface.hooks.huggingface.HfApi", autospec=True) as hf_api, + patch("datasets.load_dataset", autospec=True) as load_dataset, + ): + hf_api.return_value.dataset_info.return_value.sha = "resolved-sha" + load_dataset.return_value = dataset + operator = LoadDatasetOperator( + task_id="load", + path="org/data", + revision="release-tag", + output_uri=output, + download_mode="force_redownload", + cache_dir=str(tmp_path / "cache"), + huggingface_conn_id="huggingface_test", + ) + reference = operator.execute({}) + hf_api.return_value.dataset_info.assert_called_once_with("org/data", revision="release-tag") + load_dataset.assert_called_once_with( + "org/data", + revision="resolved-sha", + download_mode="force_redownload", + cache_dir=str(tmp_path / "cache"), + token="operator-test-token", + ) + assert reference["revision"] == "resolved-sha" + assert _load_artifact(reference).num_rows == 2 + + +def test_non_hub_load_does_not_resolve_hub_revision(tmp_path): + with patch("datasets.load_dataset", autospec=True) as load_dataset: + load_dataset.return_value = Dataset.from_dict({"x": [1]}) + operator = LoadDatasetOperator( + task_id="csv", path="csv", output_uri=(tmp_path / "csv.json").as_uri(), huggingface_conn_id=None + ) + result = operator.execute({}) + load_dataset.assert_called_once_with("csv", download_mode="reuse_dataset_if_exists") + assert result["revision"] is None + + +def test_stream_operator_lazily_transforms_filters_and_publishes_with_asset_event(): + transformed = lambda row, increment: {"value": row["value"] + increment} + predicate = lambda row: row["value"] > 1 + stream = MagicMock() + stream.map.return_value = stream + stream.filter.return_value = stream + commit = commit_info("org/streamed", message="Stream dataset from Airflow") + with patch( + "airflow.providers.huggingface.operators.datasets.HuggingFaceHook", autospec=True + ) as hook_type: + hook = hook_type.return_value + hook.get_dataset_info.return_value.sha = "pinned-stream-sha" + hook.load_dataset.return_value = stream + hook.find_commit_by_marker.return_value = None + hook.push_to_hub.return_value = commit + operator = StreamDatasetOperator( + task_id="stream", + path="org/source", + split="train", + revision="main", + repo_id="org/streamed", + config="clean", + private=True, + transform=transformed, + transform_kwargs={"increment": 1}, + predicate=predicate, + huggingface_conn_id=None, + ) + event = SimpleNamespace(extra=None) + result = operator.execute( + {"run_id": "scheduled__2026-10-31", "outlet_events": {operator.published_asset: event}} + ) + + hook.get_dataset_info.assert_called_once_with("org/source", "main") + hook.load_dataset.assert_called_once_with( + "org/source", streaming=True, split="train", revision="pinned-stream-sha" + ) + hook.push_to_hub.assert_called_once() + stream.map.assert_called_once_with(transformed, fn_kwargs={"increment": 1}) + stream.filter.assert_called_once_with(predicate) + assert hook.push_to_hub.call_args.args == (stream, "org/streamed") + push_kwargs = hook.push_to_hub.call_args.kwargs + assert push_kwargs["config_name"] == "clean" + assert push_kwargs["private"] is True + assert "huggingface-stream-publish:" in push_kwargs["commit_message"] + assert result["source_revision"] == "pinned-stream-sha" + assert result["commit_sha"] == commit.oid + assert result["streaming"] is True + assert event.extra["commit_sha"] == commit.oid + + +def test_stream_operator_recovers_prior_commit_without_reloading_stream(): + commit = commit_info("org/streamed") + with patch( + "airflow.providers.huggingface.operators.datasets.HuggingFaceHook", autospec=True + ) as hook_type: + hook = hook_type.return_value + hook.find_commit_by_marker.return_value = commit + operator = StreamDatasetOperator( + task_id="stream", + path="csv", + repo_id="org/streamed", + huggingface_conn_id=None, + ) + state = type( + "Store", + (), + { + "get": lambda self, key: True if key.startswith("stream-publication-pending:") else None, + "set": lambda self, key, value: None, + }, + )() + result = operator.execute({"run_id": "scheduled__2026-10-31", "task_state_store": state}) + + hook.load_dataset.assert_not_called() + hook.push_to_hub.assert_not_called() + assert result["commit_sha"] == commit.oid + + +def test_parquet_manifest_preserves_order_features_and_dataset_splits(tmp_path): + features = Features( + { + "z_first": Value("int64"), + "a_second": Value("string"), + "kind": ClassLabel(names=["first", "second"]), + } + ) + source = DatasetDict( + { + "train": Dataset.from_dict({"z_first": [1], "a_second": ["one"], "kind": [0]}, features=features), + "validation": Dataset.from_dict( + {"z_first": [2], "a_second": ["two"], "kind": [1]}, features=features + ), + } + ) + reference = _save(source, (tmp_path / "splits.json").as_uri(), {"revision": "abc"}) + restored = _load_artifact(reference) + assert list(restored) == ["train", "validation"] + assert restored["train"].column_names == ["z_first", "a_second", "kind"] + assert restored["validation"].features == features + assert restored["train"].features["kind"].int2str(0) == "first" + assert reference["format"] == "parquet" + assert reference["fingerprint"] + assert reference["splits"]["train"]["rows"] == 1 + + +def test_parquet_accepts_large_list_features(tmp_path): + feature = LargeList(Value("int64")) + source = Dataset.from_dict({"values": [[1, 2], [3]]}, features=Features({"values": feature})) + reference = _save(source, (tmp_path / "large-list.json").as_uri()) + restored = _load_artifact(reference) + assert restored["values"] == [[1, 2], [3]] + + +def test_single_dataset_split_name_survives_artifact_round_trip(tmp_path): + source = Dataset.from_dict({"value": [1]}, split="validation") + reference = _save(source, (tmp_path / "single-split.json").as_uri()) + restored = _load_artifact(reference) + assert restored.split == "validation" + + +def test_parquet_rejects_media_feature_types(tmp_path): + dataset = Dataset.from_dict({"image": [None]}, features=Features({"image": Image()})) + with pytest.raises(ValueError, match="unsupported Parquet feature type"): + _save(dataset, (tmp_path / "image.json").as_uri()) + + +def test_map_uses_parse_time_callable_and_persists_reference(tmp_path): + source = _save( + Dataset.from_dict({"value": [1, 2]}), + (tmp_path / "input.json").as_uri(), + {"repository_id": "org/data", "revision": "source-sha"}, + ) + source.update(repository_id="org/data", revision="source-sha") + operator = MapDatasetOperator( + task_id="map", input_uri=source, output_uri=(tmp_path / "mapped.json").as_uri(), transform=add_one + ) + reference = operator.execute({}) + assert _load_artifact(reference)["value"] == [2, 3] + assert reference["revision"] == "source-sha" + assert "transform" not in operator.template_fields + with pytest.raises(TypeError, match="callable"): + MapDatasetOperator( + task_id="bad", input_uri=source, output_uri="file:///tmp/out.json", transform="x:y" + ) + + +def test_map_transform_exception_propagates(tmp_path): + source = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + error = LookupError("transform failed") + + def fail(example): + raise error + + operator = MapDatasetOperator( + task_id="map_fail", input_uri=source, output_uri=(tmp_path / "mapped.json").as_uri(), transform=fail + ) + with pytest.raises(LookupError) as raised: + operator.execute({}) + assert raised.value is error + + +def test_fingerprint_and_split_checksum_mismatches_are_rejected(tmp_path): + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + changed_reference = dict(reference, fingerprint="wrong") + with pytest.raises(ValueError, match="artifact fingerprint"): + _load_artifact(changed_reference) + split_path = tmp_path / "input.json.split-00000.parquet" + split_path.write_bytes(b"corrupted") + with pytest.raises(ValueError, match="split fingerprint"): + _load_artifact(reference) + + +def test_push_returns_commit_info_and_emits_matching_asset_event(tmp_path, hub_dependencies): + hf_api, push_to_hub = hub_dependencies + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + commit = commit_info() + push_to_hub.return_value = commit + operator = PushToHubOperator( + task_id="publish", + input_uri=reference, + repo_id="org/data", + commit_message="publish source", + huggingface_conn_id="huggingface_test", + ) + event = type("Event", (), {"extra": None})() + + class Events: + def __getitem__(self, asset): + assert asset.uri == "hf://datasets/org/data" + return event + + result = operator.execute({"outlet_events": Events()}) + push_to_hub.assert_called_once() + kwargs = push_to_hub.call_args.kwargs + assert reference["fingerprint"] in kwargs["commit_message"] + hf_api.return_value.list_repo_commits.assert_not_called() + assert result["commit_sha"] == commit.oid + assert result["commit_url"] == commit.commit_url + assert event.extra["commit_sha"] == result["commit_sha"] + assert event.extra["commit_url"] == result["commit_url"] + + +def test_push_failure_propagates_and_emits_no_event(tmp_path, hub_dependencies): + _, push_to_hub = hub_dependencies + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + error = ValueError("publication denied") + push_to_hub.side_effect = error + operator = PushToHubOperator( + task_id="failed", input_uri=reference, repo_id="org/data", huggingface_conn_id="huggingface_test" + ) + event = type("Event", (), {"extra": None})() + + class Events: + def __getitem__(self, asset): + return event + + with pytest.raises(ValueError, match="publication denied") as raised: + operator.execute({"outlet_events": Events()}) + assert raised.value is error + assert event.extra is None + + +def test_push_requires_reference_and_validates_fingerprint(tmp_path): + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + operator = PushToHubOperator(task_id="publish", input_uri=reference["uri"], repo_id="org/data") + with pytest.raises(ValueError, match="requires an artifact reference"): + operator.execute({}) + reference["fingerprint"] = "bad" + operator = PushToHubOperator( + task_id="publish_bad", input_uri=reference, repo_id="org/data", durable=False + ) + with pytest.raises(ValueError, match="artifact fingerprint"): + operator.execute({}) + + +def test_airflow_30_path_without_state_store_and_retry_with_store(tmp_path, hub_dependencies): + hf_api, push_to_hub = hub_dependencies + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "input.json").as_uri()) + target_marker = _publication_marker( + "target", + { + "repo_id": "org/data", + "config_name": "default", + "revision": "main", + "create_pr": False, + "push_kwargs": {}, + }, + ) + marker = f"airflow-artifact-sha256:{reference['fingerprint']}:{target_marker}" + commit = git_commit_info("a" * 40, f"Publish [{marker}]") + hf_api.return_value.list_repo_commits.return_value = [commit] + push_to_hub.return_value = commit_info(oid="a" * 40) + operator = PushToHubOperator( + task_id="legacy", input_uri=reference, repo_id="org/data", huggingface_conn_id="huggingface_test" + ) + result = operator.execute({}) + assert result["commit_sha"] == "a" * 40 + hf_api.return_value.list_repo_commits.assert_not_called() + push_to_hub.assert_called_once() + store = type("Store", (), {"get": lambda self, key: result, "set": lambda self, key, value: None})() + event = type("Event", (), {"extra": None})() + + class Events: + def __getitem__(self, asset): + return event + + assert operator.execute({"task_state_store": store, "outlet_events": Events()}) == result + assert event.extra["commit_sha"] == result["commit_sha"] + + +def test_push_reuses_hub_commit_before_loading_missing_artifact(tmp_path, hub_dependencies): + hf_api, push_to_hub = hub_dependencies + reference = { + "uri": (tmp_path / "deleted.json").as_uri(), + "fingerprint": "f" * 64, + "repository_id": "org/source", + "revision": "source-sha", + } + target_marker = _publication_marker( + "target", + { + "repo_id": "org/data", + "config_name": "default", + "revision": "main", + "create_pr": False, + "push_kwargs": {}, + }, + ) + marker = f"airflow-artifact-sha256:{reference['fingerprint']}:{target_marker}" + commit = git_commit_info("b" * 40, f"Publish [{marker}]") + hf_api.return_value.list_repo_commits.return_value = [commit] + operator = PushToHubOperator( + task_id="recover_missing", + input_uri=reference, + repo_id="org/data", + huggingface_conn_id="huggingface_test", + ) + state = type( + "Store", + (), + { + "get": lambda self, key: True if key.startswith("publication-pending:") else None, + "set": lambda self, key, value: None, + }, + )() + with patch("airflow.providers.huggingface.operators.datasets._load_artifact") as load: + result = operator.execute({"task_state_store": state}) + load.assert_not_called() + push_to_hub.assert_not_called() + assert result["commit_sha"] == "b" * 40 + + +def test_state_store_write_failure_retry_does_not_push_twice(tmp_path, hub_dependencies): + hf_api, push_to_hub = hub_dependencies + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "retry_after_push.json").as_uri()) + target_marker = _publication_marker( + "target", + { + "repo_id": "org/data", + "config_name": "default", + "revision": "main", + "create_pr": False, + "push_kwargs": {}, + }, + ) + marker = f"airflow-artifact-sha256:{reference['fingerprint']}:{target_marker}" + commit = git_commit_info("a" * 40, f"Publish [{marker}]") + hf_api.return_value.list_repo_commits.return_value = [commit] + push_to_hub.return_value = commit + operator = PushToHubOperator( + task_id="retry_after_push", + input_uri=reference, + repo_id="org/data", + huggingface_conn_id="huggingface_test", + ) + store_data = {} + + class FailingStore: + def get(self, key): + return store_data.get(key) + + def set(self, key, value): + if key.startswith("publication-pending:"): + store_data[key] = value + return + raise RuntimeError("worker failed after Hub commit") + + with pytest.raises(RuntimeError, match="worker failed"): + operator.execute({"task_state_store": FailingStore()}) + + class Store: + def get(self, key): + return store_data.get(key) + + def set(self, key, value): + store_data[key] = value + + retry_result = operator.execute({"task_state_store": Store()}) + push_to_hub.assert_called_once() + assert retry_result["commit_sha"] == "a" * 40 + + +def test_push_does_not_treat_intermediate_upload_part_as_published(tmp_path, hub_dependencies): + hf_api, push_to_hub = hub_dependencies + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "partial.json").as_uri()) + target_marker = _publication_marker( + "target", + { + "repo_id": "org/data", + "config_name": "default", + "revision": "main", + "create_pr": False, + "push_kwargs": {}, + }, + ) + marker = f"airflow-artifact-sha256:{reference['fingerprint']}:{target_marker}" + hf_api.return_value.list_repo_commits.return_value = [ + git_commit_info("a" * 40, f"Publish [{marker}] (part 00000-of-00002)") + ] + pushed_commit = commit_info(oid="b" * 40) + push_to_hub.return_value = pushed_commit + + operator = PushToHubOperator( + task_id="partial_upload", + input_uri=reference, + repo_id="org/data", + huggingface_conn_id="huggingface_test", + ) + state = type( + "Store", + (), + { + "get": lambda self, key: True if key.startswith("publication-pending:") else None, + "set": lambda self, key, value: None, + }, + )() + result = operator.execute({"task_state_store": state}) + + push_to_hub.assert_called_once() + assert result["commit_sha"] == pushed_commit.oid + + +def test_push_idempotency_marker_is_scoped_to_config_revision_and_pr_target(tmp_path): + reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path / "targets.json").as_uri()) + commits = [commit_info(oid=f"{index:040x}") for index in range(1, 4)] Review Comment: `side_effect` has three commits for six configs, so the fourth `execute` raises `StopIteration` and the `split` / `data_dir` cases never get checked. `range(1, 7)`? ########## providers/huggingface/src/airflow/providers/huggingface/hooks/huggingface.py: ########## @@ -0,0 +1,144 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Connection and client helpers for Hugging Face services.""" + +from __future__ import annotations + +from typing import Any + +from huggingface_hub import HfApi +from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError + +from airflow.exceptions import AirflowNotFoundException +from airflow.providers.common.compat.sdk import BaseHook + + +class HuggingFaceHook(BaseHook): + """ + Access Hugging Face Datasets and Hub APIs using an Airflow connection. + + If ``huggingface_default`` is absent, or no connection ID is configured, + huggingface_hub may use ``HF_TOKEN`` or a cached worker login. A missing + explicitly named connection raises Airflow's connection-not-found + exception. An existing connection with no password disables SDK + authentication fallback. + """ + + conn_name_attr = "huggingface_conn_id" + default_conn_name = "huggingface_default" + conn_type = "huggingface" + hook_name = "Hugging Face" + + def __init__(self, huggingface_conn_id: str | None = default_conn_name) -> None: + super().__init__() + self.huggingface_conn_id = huggingface_conn_id + self._allow_ambient_fallback = huggingface_conn_id in (None, self.default_conn_name) + self._client: HfApi | None = None + + def get_token(self) -> str | bool | None: + if not self.huggingface_conn_id: + return None + try: + connection = self.get_connection(self.huggingface_conn_id) + except AirflowNotFoundException: + if self._allow_ambient_fallback: + return None + raise + return connection.password or False + + def get_conn(self) -> HfApi: + if self._client is None: + self._client = HfApi(token=self.get_token()) + return self._client + + def test_connection(self) -> tuple[bool, str]: + try: + self.get_conn().whoami() + return True, "Connection successfully tested" + except Exception as exc: + return False, str(exc) + + def load_dataset(self, path: str, **kwargs: Any) -> Any: + from datasets import load_dataset + + token = self.get_token() + if token is not None and "token" not in kwargs: + kwargs["token"] = token + return load_dataset(path, **kwargs) + + def get_dataset_info(self, repo_id: str, revision: str | None = None) -> Any: + return self.get_conn().dataset_info(repo_id, revision=revision) + + @staticmethod + def _commit_is_complete(commit: Any) -> bool: + """Reject intermediate commits from a Hub upload split into parts.""" + import re + + title = getattr(commit, "title", "") or "" + return re.search(r"\(part\b.*\)\s*$", title, re.IGNORECASE) is None + + def find_commit_by_marker( + self, + repo_id: str, + marker: str, + *, + revision: str | None = None, + include_pull_requests: bool = False, + ) -> Any | None: + """Find the latest complete commit matching a marker on the requested ref.""" + api = self.get_conn() + revisions = [revision] if revision else ["main"] + if include_pull_requests: + try: + refs = api.list_repo_refs(repo_id, repo_type="dataset", include_pull_requests=True) + except RepositoryNotFoundError: + return None + revisions.extend(ref.ref for ref in getattr(refs, "pull_requests", []) or []) + + seen: set[str] = set() + for ref in revisions: + if ref in seen: + continue + seen.add(ref) + try: + commits = api.list_repo_commits(repo_id, repo_type="dataset", revision=ref) + except RepositoryNotFoundError: + return None + except RevisionNotFoundError: + continue + # Hub returns commits oldest first. A historical match is stale if + # a later publication has advanced this ref since the push attempt. + commit = commits[-1] if commits else None Review Comment: `list_repo_commits` returns newest first (its docstring says "last commit first", and its own example takes `[-1]` as the initial commit), so `commits[-1]` is the repo's system `initial commit` and never carries the marker. A retry after the worker dies between `push_to_hub` and `state.set(state_key, ...)` pushes again, and the post-push lookups in both operators can't find anything either. The hook tests pass because they build the list oldest first. `commits[0]`, with the fixtures in Hub order? ########## providers/huggingface/docs/index.rst: ########## @@ -0,0 +1,164 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Apache Airflow Hugging Face Provider +==================================== + +The Hugging Face provider connects Airflow tasks to Hugging Face Datasets and +Hub repositories. It supplies a shared ``HuggingFaceHook`` and four operators +to load a dataset, map a Python callable, stream directly to the Hub, and +publish a durable artifact. Durable data moves between tasks as Parquet objects +and a small JSON manifest, so tasks can run on separate workers when their +artifact store is shared. + +Supported capabilities +---------------------- + +* Airflow connection based Hub authentication, with Hub SDK fallback when the + default connection is absent. +* Revision pinning and source provenance for Hub dataset loads. +* Split aware Parquet handoff with features metadata and content fingerprints. +* Synchronous publication of materialized or streamed datasets with the Hub + commit SHA and URL in results. +* ``hf://datasets/<repo_id>`` assets and OpenLineage conversion. + +``StreamDatasetOperator`` lazily transforms and filters an ``IterableDataset`` +and publishes it directly to the Hub. Datasets uploads Parquet shards +progressively; the operator does not retain all rows or create an intermediate +artifact. Standalone filtering and split assembly operators are not included +in the v1 public API. + +Installation +------------ + +Install into an Airflow environment with Python 3.11 through 3.14: + +.. code-block:: bash + + pip install apache-airflow-providers-huggingface + +The provider requires ``apache-airflow>=3.0,<4``, +``apache-airflow-providers-common-compat>=1.8.0``, ``datasets>=5.0.0``, and +``huggingface-hub>=0.24``. For ``s3://`` and ``gs://`` artifacts, install the +corresponding storage provider and filesystem dependency. Select +the matching Airflow connection in the URI's user-info section, for example +``s3://aws_default@my-bucket/path/data.json`` or +``gs://google_cloud_default@my-bucket/path/data.json``. Without a connection +ID in the URI, the filesystem uses its default credential chain; configuring +an Airflow connection alone does not select it. If an operator's +``storage_conn_id`` is set, do not also put a connection ID in the URI: which +one takes precedence differs between supported Airflow versions. + +Guides +------ + +.. toctree:: + :maxdepth: 2 + + connections Review Comment: Sphinx doesn't resolve a bare `connections` to `connections/index`, so this toctree points at a missing doc and both pages end up orphaned, and the docs build turns those warnings into errors. `` :doc:`artifacts` `` in operators/index.rst has the same problem, it resolves to `operators/artifacts`. `Connection types <connections/index>` and `Operators <operators/index>` like amazon does, and `` :doc:`/artifacts` ``? ########## pyproject.toml: ########## @@ -259,6 +259,9 @@ apache-airflow = "airflow.__main__:main" "hashicorp" = [ "apache-airflow-providers-hashicorp>=4.0.0" ] +"huggingface" = [ Review Comment: `prek run --all-files` still fails at this head, though most of it is auto-fixed: ruff and ruff-format, the provider order here (huggingface goes after http, in this extra, `all` and the other lists), license, final-newline and trailing-space fixes in `NOTICE` and the docs, the missing `test_run = get_test_run(dag)` in `example_read_only.py`, and five `release-management` breeze images now that it's `ready`. The one it can't fix is `validate-operators-init`: `MapDatasetOperator` sets `self.transform_kwargs = {} if transform_kwargs is None else transform_kwargs`, which counts as logic on a template field, so that default has to move to `execute`. Could you run it and push what it rewrites? ########## providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py: ########## @@ -0,0 +1,796 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Operators to load, transform, and publish Hugging Face datasets.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import json +import marshal +import tempfile +from collections.abc import Callable +from contextlib import suppress +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from datasets import Dataset, DatasetDict + +from airflow.providers.common.compat.sdk import Asset, BaseOperator, Context, ObjectStoragePath +from airflow.providers.huggingface.hooks.huggingface import HuggingFaceHook + + +def _path(uri: str, *, conn_id: str | None = None) -> ObjectStoragePath: + """Resolve a URI through Airflow's configured object-storage filesystems.""" + return ObjectStoragePath(uri, conn_id=conn_id) + + +def _reference(reference: str | dict[str, Any]) -> str: + if isinstance(reference, dict): + uri = reference.get("uri") + if not uri: + raise ValueError("Artifact reference must contain a non-empty 'uri'.") + return uri + if not isinstance(reference, str) or not reference: + raise ValueError("Artifact reference must be a URI or contain a 'uri'.") + return reference + + +class _HashingWriter: + def __init__(self, stream: Any) -> None: + self._stream = stream + self._digest = hashlib.sha256() + + def write(self, data: bytes) -> int: + written = self._stream.write(data) + self._digest.update(data[:written]) + return written + + def get_hexdigest(self) -> str: + return self._digest.hexdigest() + + def __getattr__(self, name: str) -> Any: + return getattr(self._stream, name) + + +def _validate_feature(feature: Any, field_name: str) -> None: + """Allow Parquet-safe scalar, class-label, and nested sequence features.""" + from datasets import ClassLabel, LargeList, Sequence, Value + + if isinstance(feature, (Value, ClassLabel)): + return + if isinstance(feature, (Sequence, LargeList)): + _validate_feature(feature.feature, field_name) + return + if isinstance(feature, dict): + for key, value in feature.items(): + _validate_feature(value, f"{field_name}.{key}") + return + if isinstance(feature, (list, tuple)): + for value in feature: + _validate_feature(value, field_name) + return + raise ValueError( + f"unsupported Parquet feature type for {field_name!r}: {type(feature).__name__}; " + "supported types are Value, ClassLabel, and nested Sequence or LargeList features" + ) + + +def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str, Dataset]]: + from datasets import Dataset, DatasetDict + + if isinstance(dataset, DatasetDict): + return list(dataset.items()) + if isinstance(dataset, Dataset): + return [("dataset", dataset)] + raise TypeError("Expected a datasets.Dataset or datasets.DatasetDict.") + + +def _save( + dataset: Dataset | DatasetDict, + uri: str, + metadata: dict[str, Any] | None = None, + *, + storage_conn_id: str | None = None, +) -> dict[str, Any]: + """Write split Parquet objects and a JSON manifest through ObjectStoragePath.""" + import pyarrow.parquet as parquet + + parts = _dataset_parts(dataset) + written: list[ObjectStoragePath] = [] + manifest_path = _path(uri, conn_id=storage_conn_id) + split_manifest: dict[str, Any] = {} + try: + for index, (name, part) in enumerate(parts): + for field_name, feature in part.features.items(): + _validate_feature(feature, field_name) + split_uri = f"{uri}.split-{index:05d}.parquet" + split_path = _path(split_uri, conn_id=storage_conn_id) + split_path.parent.mkdir(parents=True, exist_ok=True) + with split_path.open("wb") as stream: + hashing_stream = _HashingWriter(stream) + parquet.write_table(part.data.table, hashing_stream) + written.append(split_path) + split_manifest[name] = { + "uri": split_uri, + "rows": len(part), + "features": part.features.to_dict(), + "feature_order": list(part.features), + "split": ( + str(part.split) + if getattr(part, "split", None) is not None + else (name if name != "dataset" else None) + ), + "sha256": hashing_stream.get_hexdigest(), + } + + manifest = { + "format": "parquet", + "version": 1, + "splits": split_manifest, + "metadata": metadata or {}, + } + manifest_bytes = json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest_path.parent.mkdir(parents=True, exist_ok=True) + with manifest_path.open("wb") as stream: + stream.write(manifest_bytes) + return { + "uri": uri, + "format": "parquet", + "splits": { + name: { + "rows": part_info["rows"], + "features": part_info["features"], + "split": part_info["split"], + } + for name, part_info in split_manifest.items() + }, + **(metadata or {}), + "fingerprint": hashlib.sha256(manifest_bytes).hexdigest(), + } + except Exception: + # Remove partial split outputs. A failed write must not leave an artifact + # reference that a downstream task could mistake for a complete dataset. + for path in written: + with suppress(OSError, NotImplementedError): + path.unlink(missing_ok=True) + with suppress(OSError, NotImplementedError): + manifest_path.unlink(missing_ok=True) + raise + + +def _read_artifact_manifest( + reference: str | dict[str, Any], *, storage_conn_id: str | None = None +) -> tuple[str, dict[str, Any]]: + uri = _reference(reference) + manifest_path = _path(uri, conn_id=storage_conn_id) + with manifest_path.open("rb") as stream: + manifest_bytes = stream.read() + if isinstance(reference, dict) and reference.get("fingerprint"): + actual = hashlib.sha256(manifest_bytes).hexdigest() + if actual != reference["fingerprint"]: + raise ValueError("Dataset artifact fingerprint does not match its reference") + + manifest = json.loads(manifest_bytes) + if manifest.get("format") != "parquet" or manifest.get("version") != 1: + raise ValueError(f"Unsupported dataset artifact manifest: {uri}") + return uri, manifest + + +def _load_artifact( + reference: str | dict[str, Any], + *, + cache_dir: str | None = None, + manifest: dict[str, Any] | None = None, + storage_conn_id: str | None = None, +) -> Dataset | DatasetDict: + from datasets import Dataset, DatasetDict, Features + + uri, manifest = ( + _read_artifact_manifest(reference, storage_conn_id=storage_conn_id) + if manifest is None + else (_reference(reference), manifest) + ) + resolved_cache_dir = cache_dir or str(Path(tempfile.gettempdir()) / "airflow-hf-datasets-cache") + Path(resolved_cache_dir).mkdir(parents=True, exist_ok=True) + parts: dict[str, Dataset] = {} + with tempfile.TemporaryDirectory(prefix="airflow-hf-parquet-") as local_dir: + for index, (name, item) in enumerate(manifest["splits"].items()): + split_uri = item["uri"] + local_path = Path(local_dir) / f"split-{index:05d}.parquet" + digest = hashlib.sha256() + with ( + _path(split_uri, conn_id=storage_conn_id).open("rb") as source, + local_path.open("wb") as target, + ): + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + target.write(chunk) + if digest.hexdigest() != item.get("sha256"): + raise ValueError(f"Dataset artifact split fingerprint mismatch: {split_uri}") + + feature_data = item["features"] + feature_order = item.get("feature_order", list(feature_data)) + features = Features.from_dict({key: feature_data[key] for key in feature_order}) + if item.get("rows") == 0: + restored = Dataset.from_dict({key: [] for key in features}, features=features) + else: + restored = Dataset.from_parquet( + str(local_path), + features=features, + cache_dir=resolved_cache_dir, + ) + # Parquet readers may return fields in physical/file order; restore the + # manifest's declared feature order for stable downstream transforms. + parts[name] = restored.select_columns(list(features)) + return parts["dataset"] if list(parts) == ["dataset"] else DatasetDict(parts) + + +def _state_key(prefix: str, identity: str) -> str: + """Task state stores already scope keys to a DAG run and task instance.""" + return f"{prefix}:{identity}" + + +def _publication_marker(prefix: str, identity: dict[str, Any]) -> str: + payload = json.dumps(identity, sort_keys=True, separators=(",", ":"), default=str).encode() + return f"{prefix}:{hashlib.sha256(payload).hexdigest()}" + + +def _callable_identity(function: Callable[..., Any] | None) -> dict[str, str] | None: + if function is None: + return None + code = getattr(function, "__code__", None) + if code is None: + code = getattr(type(function).__call__, "__code__", None) + return { + "module": getattr(function, "__module__", type(function).__module__), + "qualname": getattr(function, "__qualname__", type(function).__qualname__), + "code_hash": hashlib.sha256(marshal.dumps(code)).hexdigest() if code is not None else "", + } + + +def _warn_missing_task_state(operator: BaseOperator, state: Any) -> None: + if getattr(operator, "durable", False) and state is None: + operator.log.warning( + "durable=True requires Airflow 3.3 task state store support; revision pinning and " + "task-state result caching are unavailable for this task run" + ) + + +def _commit_details(commit: Any, repo_id: str, hook: HuggingFaceHook) -> tuple[str, str]: + """Normalize CommitInfo and GitCommitInfo returned by Hub APIs.""" + commit_sha = getattr(commit, "oid", None) or getattr(commit, "commit_id", None) + if not commit_sha: + raise ValueError(f"Hub publication for {repo_id!r} returned no commit identifier") + commit_url = getattr(commit, "commit_url", None) + if not commit_url: + endpoint = getattr(hook.get_conn(), "endpoint", "https://huggingface.co";).rstrip("/") + commit_url = f"{endpoint}/datasets/{repo_id}/commit/{commit_sha}" + return commit_sha, commit_url + + +def _is_hub_repo(path: str) -> bool: + """Treat non-local paths as Hub IDs unless they name a packaged builder.""" + if Path(path).exists(): + return False + if "/" in path: + return True + try: + return importlib.util.find_spec(f"datasets.packaged_modules.{path}") is None + except (ImportError, ModuleNotFoundError, ValueError): + return True + + +def _publication_outlets( + repo_id: str, outlets: Any, push_kwargs: dict[str, Any] +) -> tuple[Asset | None, list[Any]]: + declared_outlets = list(outlets) if isinstance(outlets, (list, tuple)) else ([outlets] if outlets else []) + if push_kwargs.get("create_pr") or push_kwargs.get("revision", "main") not in (None, "main"): + return None, declared_outlets + + asset_uri = f"hf://datasets/{repo_id}" + existing_asset = next( + (outlet for outlet in declared_outlets if getattr(outlet, "uri", None) == asset_uri), None + ) + asset = existing_asset or Asset(uri=asset_uri) + if existing_asset is None: + declared_outlets.append(asset) + return asset, declared_outlets + + +class LoadDatasetOperator(BaseOperator): + """Load a dataset revision and persist it as a Parquet artifact.""" + + template_fields = ( + "path", + "name", + "split", + "revision", + "cache_dir", + "output_uri", + "storage_conn_id", + "huggingface_conn_id", + ) + + def __init__( + self, + *, + path: str, + output_uri: str, + name: str | None = None, + split: str | None = None, + revision: str | None = None, + cache_dir: str | None = None, + storage_conn_id: str | None = None, + download_mode: str = "reuse_dataset_if_exists", + huggingface_conn_id: str | None = "huggingface_default", + load_kwargs: dict[str, Any] | None = None, + durable: bool = True, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.output_uri = output_uri + self.name = name + self.split = split + self.revision = revision + self.cache_dir = cache_dir + self.download_mode = download_mode + self.storage_conn_id = storage_conn_id + self.huggingface_conn_id = huggingface_conn_id + self.load_kwargs = load_kwargs or {} + self.durable = durable + + def execute(self, context: Context) -> dict[str, Any]: + hook = HuggingFaceHook(self.huggingface_conn_id) + requested_revision = self.revision + is_hub_repo = _is_hub_repo(self.path) + state = context.get("task_state_store") if self.durable else None + _warn_missing_task_state(self, state) + state_identity = json.dumps([self.path, requested_revision], separators=(",", ":")) Review Comment: The raw path goes into this key, so any `org/name` source puts a `/` in the task state key. On 3.3.0 the SDK client doesn't quote the key and the route is `/{task_instance_id}/{key}`, so `state.set` gets a 404 and Load fails with the default `durable=True` (Stream too, line 477). 3.3.1 added `quote()` and `{key:path}`. Could the identity be hashed, like `_publication_marker` does? That also keeps it under the 512-char key column. ########## providers/huggingface/docs/operators/index.rst: ########## @@ -0,0 +1,212 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +Operators +========= + +The initial public operator API is ``LoadDatasetOperator``, +``MapDatasetOperator``, ``StreamDatasetOperator``, and +``PushToHubOperator``. The operators exchange the manifest and split Parquet +objects described in :doc:`artifacts`. + +LoadDatasetOperator +------------------- + +``LoadDatasetOperator`` accepts ``path``, ``output_uri``, optional ``name`` +(dataset configuration), ``split``, ``revision``, ``cache_dir``, +``download_mode``, ``load_kwargs``, ``storage_conn_id``, and ``durable`` +(default ``True``). ``storage_conn_id`` selects the Airflow object-storage +connection for ``output_uri``. If omitted, a connection ID in the URI's +user-info section or the backend's default credentials are used. Set only one +of ``storage_conn_id`` and a URI user-info connection ID, because precedence +differs between supported Airflow versions. It +persists the result and returns the +artifact URI plus source ID, resolved revision, configuration, split, row +counts, and feature descriptions. It resolves a Hub source to a commit before +loading; pin a commit SHA in ``revision`` for reproducible inputs. Unnamespaced +IDs such as ``imdb`` are treated as Hub repositories unless they name an +existing local path or a packaged Datasets builder such as ``csv`` or +``parquet``. + +With ``durable=True``, the operator stores the resolved source revision in +Airflow's task state store so retries use the same Hub commit. Set it to +``False`` to resolve the revision again on each attempt without task-state +reuse. The task state store is available in Airflow 3.3 and later; on earlier +versions, ``durable=True`` logs a warning and cannot pin the revision across +retries. + +The default ``reuse_dataset_if_exists`` mode reuses Datasets' valid cache. +``force_redownload`` forces a new download. ``cache_dir`` isolates cache data; +it is separate from the output artifact destination. Keeping a stable cache +directory allows retries and cleared task instances to reuse downloaded data. + +MapDatasetOperator +------------------ + +Required parameters are ``input_uri``, ``output_uri``, and ``transform``. +``transform`` must be a Python callable available when the DAG is parsed and +when the task worker executes. The operator does not template or import the +callable from a string. ``transform_kwargs`` are passed as Datasets ``fn_kwargs``; +``map_kwargs`` are passed to ``Dataset.map`` or ``DatasetDict.map``. The output +preserves compatible provenance and records a new artifact fingerprint. +``storage_conn_id`` selects the connection used for both input and output +artifact paths. Use it or connection IDs in the artifact URIs' user-info +sections, but not both, because precedence differs between supported Airflow +versions. + +The system example demonstrates the durable artifact flow and its task handoff: + +.. exampleinclude:: /../../huggingface/tests/system/huggingface/example_read_only.py + :language: python + :start-after: [START huggingface_load_and_map_example] + :end-before: [END huggingface_load_and_map_example] + +StreamDatasetOperator +--------------------- + +Hugging Face Datasets ``streaming=True`` returns an ``IterableDataset`` (or an +iterable of splits) that reads examples incrementally. It avoids materializing +the full source dataset and is intended for sequential iteration; it does not +provide the usual in-memory ``Dataset`` behavior such as random access. + +``StreamDatasetOperator`` loads with +``streaming=True``, applies the optional ``transform`` with +``transform_kwargs``, then applies the optional ``predicate`` using Datasets' +lazy iterable operations. It passes the resulting ``IterableDataset`` or +``IterableDatasetDict`` directly to ``push_to_hub`` in the same task. Datasets +progressively writes Parquet shards to the Hub, so rows are not accumulated in +a Python list, stored in XCom, or written to a local artifact first. + +The operator accepts ``path``, ``repo_id``, optional ``name``, ``split``, +``revision``, ``config``, ``private``, ``commit_message``, +``huggingface_conn_id``, ``push_kwargs``, and ``durable`` (default ``True``). +With durability enabled on Airflow 3.3 and later, it pins the resolved Hub +source revision and reuses a completed push from task state. It searches Hub +commit markers only after task state records that a push may have started. +Markers include the Dag, task, run, map index, source, and transform identity. +Set ``durable=False`` to disable those reuse checks. It returns the Hub repository, +commit SHA and URL, source revision, configuration, split, and +``streaming=True``. After a successful push it emits the matching +``hf://datasets/<repo_id>`` asset event. A run marker in the commit message +identifies the source and target for this DAG run. A retry searches the target +branch or pull request refs only when task state says a prior attempt may have +started a push. Recovery accepts the latest commit on the target ref only when +it carries this marker and has no multipart suffix. Intermediate upload parts +and older historical commits do not count as completed publication. +Mapping is supported when ``repo_id`` is fixed in ``partial()``; the mapped +Dag then declares the same repository-level asset as an ordinary task. A +``repo_id`` supplied through ``expand()`` is not supported because the asset +must be known when the Dag is serialized. Pull request publications and pushes +to revisions other than ``main`` do not emit this repo-level asset event, +because it would schedule consumers against ``main`` rather than the published +revision. + +Execution proceeds in this order: + +#. **Pin the source.** For a Hub dataset, the operator resolves the requested + revision to a commit SHA. With ``durable=True``, it stores that SHA in the + task state store under a key for the source path, requested revision, + configuration, and split. A retry reuses the SHA, so a moving branch such as + ``main`` does not silently change between attempts. Non-Hub paths are passed + to Datasets without Hub revision lookup. +#. **Check for a completed publication.** The operator builds a marker from the + Dag, task, run, map index, source and pinned revision, configuration and split, + transform and predicate identities, destination repository, target + config/revision, and pull-request setting. On Airflow 3.3 and later, it checks + task state first. It searches Hub refs only when task state records that a + prior attempt may have started the push. A matching final commit recovers an + interrupted publication. Clearing the task removes that marker state, so a Review Comment: Clear keeps task state on 3.3, and a Dag-run clear keeps it on 3.4 too (only the task-instance clear with `keep_task_state` off discards it, from 3.4). In those cases a cleared task finds the stored result, returns the old commit and emits the asset event again without pushing anything, so this sentence needs the version caveat. ########## uv.lock: ########## @@ -13235,11 +13305,16 @@ wheels = [ [[package]] name = "fsspec" -version = "2026.9.0" +version = "2026.7.0" Review Comment: The PR description doesn't mention this yet: the lock moves fsspec 2026.9.0 -> 2026.7.0, s3fs and gcsfs back too, and PyAthena 4.2.1 -> 3.40.0 for every provider, because datasets 5.1.0 caps `fsspec<=2026.7.0`. datasets 5.2.0 (released 2026-11-09) raises the cap to 2026.9.0 and PyAthena 4.2.1 needs `fsspec>=2026.9.0`, so relocking once 5.2.0 clears the 4-day `exclude-newer` window should bring them all back. Can you relock then, or call out the downgrade in the description until it does? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
