kaxil commented on code in PR #76317:
URL: https://github.com/apache/airflow/pull/76317#discussion_r4504580614
##########
pyproject.toml:
##########
@@ -259,6 +259,9 @@ apache-airflow = "airflow.__main__:main"
"hashicorp" = [
"apache-airflow-providers-hashicorp>=4.0.0"
]
+"huggingface" = [
Review Comment:
provider.yaml says `not-ready` now, but this `huggingface` extra and the
`all` entry on line 484 are still here, so `apache-airflow[all]` still asks for
a 0.1.0 that isn't on PyPI. `prek run update-pyproject-toml --all-files`
removes both, since the generator skips not-ready providers. The `huggingface`
row in `extra-packages-ref.rst` would go too.
##########
providers/huggingface/docs/index.rst:
##########
@@ -0,0 +1,123 @@
+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 or later, ``datasets>=4.0.0,<5``,
+and ``huggingface-hub>=0.24,<2``. 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.
+
+Guides
+------
+
+.. toctree::
+ :maxdepth: 2
+
+ connections
+ operators
+ artifacts
+ security
+ changelog
+
+The system-test DAGs under ``tests/system/huggingface/`` also serve as runnable
+examples and are collected with pytest through Airflow's ``get_test_run(dag)``
+system-test helper. Run them with the repository's system-test pytest setup.
+The read-only workflow uses a pinned public dataset. The publishing workflow
+uses the ``airflow-hf`` organization by default and requires a write-scoped
+``huggingface_default`` Airflow connection. Set ``HF_TEST_HUB_ORG`` or
+``HF_TEST_HUB_CONN_ID`` to override those defaults. ``HF_TEST_WRITE_REPO_ID``
+and ``HF_TEST_DURABLE_REPO_ID`` can override the base repository names. Each
+run appends ``SYSTEM_TESTS_ENV_ID`` (or a generated UUID) to its repository
+names to avoid collisions. Cleanup tasks use ``ALL_DONE`` so they are attempted
+after failures. On Airflow 3.3 and later, the publishing DAG also checks retry
+reuse of a successful Hub commit.
+
+.. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE
OVERWRITTEN AT RELEASE TIME!
+
+
+.. toctree::
+ :hidden:
+ :maxdepth: 1
+ :caption: Commits
+
+ Detailed list of commits <commits>
+
+
+apache-airflow-providers-huggingface package
+------------------------------------------------------
+
+Load, transform, and publish Hugging Face Datasets.
+
+Release: 0.1.0
+
+Provider package
+----------------
+
+This package is for the ``huggingface`` provider.
+All classes for this package are included in the
``airflow.providers.huggingface`` python package.
+
+Installation
+------------
+
+You can install this package on top of an existing Airflow installation via
+``pip install apache-airflow-providers-huggingface``.
+For the minimum Airflow version supported, see ``Requirements`` below.
+
+Requirements
+------------
+
+The minimum Apache Airflow version supported by this provider distribution is
``3.0``.
+
+=================== ==================
+PIP package Version required
+=================== ==================
+``apache-airflow`` ``>=3.0,<4``
+``datasets`` ``>=3.4.0,<5``
Review Comment:
This table and the one in README.rst still say `datasets>=3.4.0` and don't
list `apache-airflow-providers-common-compat>=1.8.0`. README.rst also still
documents the `[common.compat]` extra, which no longer exists. Both are
generated from pyproject.toml, so regenerating them should sort it.
##########
providers/huggingface/tests/unit/huggingface/hooks/test_huggingface_hook.py:
##########
@@ -0,0 +1,160 @@
+# 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-backed authentication tests for HuggingFaceHook."""
Review Comment:
`test_project_structure.py` expects each test file to mirror its module
name, so it looks for `hooks/test_huggingface.py`, `operators/test_datasets.py`
and `tests/unit/huggingface/test_asset.py`, and fails on the current names. Can
the three files be renamed, with `test_asset.py` moved up a level?
##########
providers/huggingface/provider.yaml:
##########
@@ -0,0 +1,56 @@
+# 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.
+
+package-name: apache-airflow-providers-huggingface
+name: Hugging Face
+description: Load, transform, and publish Hugging Face Datasets.
+state: not-ready
+source-date-epoch: 0
+lifecycle: incubation
+versions:
+ - 0.1.0
+integrations:
+ - integration-name: Hugging Face Datasets
Review Comment:
`run_provider_yaml_files_check.py` compares the docs and logos on disk with
provider.yaml and only skips suspended providers, so `operators.rst` needs a
`how-to-guide` entry and `HuggingFace.png` a `logo` entry under this
integration. The PNG is also 1.3 MB (`HTTP.png` is 18 KB) and ships in the
sdist through `docs/`. Could it be shrunk?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Sequence,
Value
+
+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
+
+
+def _sha256_path(uri: str, *, conn_id: str | None = None) -> str:
+ digest = hashlib.sha256()
+ with _path(uri, conn_id=conn_id).open("rb") as stream:
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _validate_feature(feature: Any, field_name: str) -> None:
+ """Allow Parquet-safe scalar, class-label, and nested sequence features."""
+ if isinstance(feature, (Value, ClassLabel)):
+ return
+ if isinstance(feature, Sequence):
+ _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 features"
+ )
+
+
+def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str,
Dataset]]:
+ 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."""
+ 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:
+ parquet.write_table(part.data.table, 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": _sha256_path(split_uri, conn_id=storage_conn_id),
+ }
+
+ 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:
+ try:
+ path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ try:
+ manifest_path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ 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:
+ 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})
+ raw_split = item.get("split") or (name if name != "dataset" else
None)
+ split_name = raw_split.split("[")[0] if raw_split and "[" in
raw_split else raw_split
+ restored = Dataset.from_parquet(
+ str(local_path),
+ features=features,
+ split=split_name,
+ 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=(",",
":")).encode()
+ return f"{prefix}:{hashlib.sha256(payload).hexdigest()}"
+
+
+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:
+ """Match a Hub dataset repo ID without treating an existing local path as
remote."""
+ return path.count("/") == 1 and not path.startswith((".", "/")) and not
Path(path).exists()
+
+
+class LoadDatasetOperator(BaseOperator):
+ """Load a dataset revision and persist it as a Parquet artifact."""
+
+ template_fields = ("path", "name", "split", "revision", "cache_dir",
"output_uri")
+
+ 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, self.output_uri, self.name, self.split = path, output_uri,
name, split
+ self.revision, self.cache_dir, self.download_mode = revision,
cache_dir, download_mode
+ self.storage_conn_id = storage_conn_id
+ self.huggingface_conn_id, self.load_kwargs = huggingface_conn_id,
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
+ state_identity = json.dumps([self.path, requested_revision],
separators=(",", ":"))
+ state_key = _state_key("load-revision", state_identity)
+ revision = state.get(state_key) if state is not None else None
+ if revision is None:
+ revision = requested_revision
+ if is_hub_repo:
+ revision = hook.get_dataset_info(self.path,
requested_revision).sha
+ if state is not None and revision is not None:
+ state.set(state_key, revision)
+ elif state is not None:
+ self.log.info("Reusing pinned Hugging Face dataset revision %s for
%s", revision, self.path)
+
+ kwargs = dict(self.load_kwargs)
+ kwargs["download_mode"] = self.download_mode
+ if self.name is not None:
+ kwargs["name"] = self.name
+ if self.split is not None:
+ kwargs["split"] = self.split
+ if revision is not None:
+ kwargs["revision"] = revision
+ if self.cache_dir is not None:
+ kwargs["cache_dir"] = self.cache_dir
+ dataset = hook.load_dataset(self.path, **kwargs)
+ metadata = {
+ "repository_id": self.path if is_hub_repo else None,
+ "revision": revision,
+ "configuration": self.name,
+ "split": self.split,
+ }
+ return _save(dataset, self.output_uri, metadata,
storage_conn_id=self.storage_conn_id)
+
+
+class StreamDatasetOperator(BaseOperator):
+ """Transform a streamed dataset lazily and publish it directly to the
Hub."""
+
+ template_fields = ("path", "name", "split", "revision", "config",
"commit_message")
+
+ def __init__(
+ self,
+ *,
+ path: str,
+ repo_id: str,
+ transform: Callable[..., dict[str, Any]] | None = None,
+ predicate: Callable[[dict[str, Any]], bool] | None = None,
+ name: str | None = None,
+ split: str | None = None,
+ revision: str | None = None,
+ config: str | None = None,
+ private: bool | None = None,
+ commit_message: str | None = None,
+ transform_kwargs: dict[str, Any] | None = None,
+ push_kwargs: dict[str, Any] | None = None,
+ huggingface_conn_id: str | None = "huggingface_default",
+ durable: bool = True,
+ **kwargs: Any,
+ ) -> None:
+ if "{{" in repo_id:
+ raise ValueError("repo_id cannot be templated: it defines the
Asset that downstream DAGs schedule on.")
+ if transform is not None and not callable(transform):
+ raise TypeError("transform must be a callable or None")
+ if predicate is not None and not callable(predicate):
+ raise TypeError("predicate must be a callable or None")
+ self.path, self.repo_id = path, repo_id
+ self.transform, self.predicate = transform, predicate
+ self.name, self.split, self.revision = name, split, revision
+ self.config, self.private, self.commit_message = config, private,
commit_message
+ self.transform_kwargs = transform_kwargs or {}
+ self.push_kwargs = push_kwargs or {}
+ self.huggingface_conn_id = huggingface_conn_id
+ self.durable = durable
+ self.published_asset = Asset(uri=f"hf://datasets/{repo_id}")
+ user_outlets = kwargs.pop("outlets", None)
+ if user_outlets is None:
+ user_outlets = []
+ elif not isinstance(user_outlets, (list, tuple)):
+ user_outlets = [user_outlets]
+ super().__init__(outlets=[*user_outlets, self.published_asset],
**kwargs)
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ hook = HuggingFaceHook(self.huggingface_conn_id)
+ is_hub_repo = _is_hub_repo(self.path)
+ requested_revision = self.revision
+ state = context.get("task_state_store") if self.durable else None
+ revision_identity = json.dumps(
+ [self.path, requested_revision, self.name, self.split],
separators=(",", ":")
+ )
+ revision_key = _state_key("stream-revision", revision_identity)
+ revision = state.get(revision_key) if state is not None else None
+ if revision is None:
+ revision = requested_revision
+ if is_hub_repo:
+ revision = hook.get_dataset_info(self.path,
requested_revision).sha
+ if state is not None and revision is not None:
+ state.set(revision_key, revision)
+ elif state is not None:
+ self.log.info("Reusing pinned Hugging Face streaming revision %s
for %s", revision, self.path)
+
+ run_id = context.get("run_id", "manual")
+ push_kwargs = dict(self.push_kwargs)
+ if self.config is not None:
+ push_kwargs["config_name"] = self.config
+ if self.private is not None:
+ push_kwargs["private"] = self.private
+ marker = _publication_marker(
+ "huggingface-stream-publish",
+ {
+ "task_id": self.task_id,
+ "run_id": run_id,
+ "source": self.path,
+ "source_revision": revision,
+ "name": self.name,
+ "split": self.split,
+ "repo_id": self.repo_id,
+ "config_name": push_kwargs.get("config_name", "default"),
+ "revision": push_kwargs.get("revision", "main"),
+ "create_pr": bool(push_kwargs.get("create_pr", False)),
+ },
+ )
+ state_key = _state_key("stream-publication", marker)
+ if state is not None:
+ prior_result = state.get(state_key)
+ if prior_result:
+ self.log.info("Reusing stored streaming publication result for
%s", self.repo_id)
+ self._emit_asset_event(context, prior_result)
+ return prior_result
+
+ target_revision = push_kwargs.get("revision")
+ include_pull_requests = bool(push_kwargs.get("create_pr", False))
+ commit = None
+ if self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=include_pull_requests,
+ )
+ if commit is not None:
+ self.log.info("Reusing existing completed streaming
publication to %s", self.repo_id)
+
+ if commit is None:
+ load_kwargs: dict[str, Any] = {"streaming": True}
+ if self.name is not None:
+ load_kwargs["name"] = self.name
+ if self.split is not None:
+ load_kwargs["split"] = self.split
+ if revision is not None:
+ load_kwargs["revision"] = revision
+ stream = hook.load_dataset(self.path, **load_kwargs)
+
+ if self.transform is not None:
+ stream = stream.map(self.transform,
fn_kwargs=self.transform_kwargs)
+ if self.predicate is not None:
+ stream = stream.filter(self.predicate)
+
+ message = self.commit_message or "Stream dataset from Airflow"
+ push_kwargs["commit_message"] = f"{message} [{marker}]"
+ commit = hook.push_to_hub(stream, self.repo_id, **push_kwargs)
+ # A retry can recover the completed Hub commit if the worker exited
+ # after upload but before Airflow recorded the task result.
+ if commit is None:
+ if self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=include_pull_requests,
+ )
+
+ if commit is None:
+ raise RuntimeError(
+ f"Hub push to {self.repo_id!r} completed without returning or
exposing a commit for marker {marker}"
+ )
+ commit_sha, commit_url = _commit_details(commit, self.repo_id, hook)
+ result = {
+ "repository_id": self.repo_id,
+ "commit_sha": commit_sha,
+ "commit_url": commit_url,
+ "url": f"https://huggingface.co/datasets/{self.repo_id}";,
+ "source_repository_id": self.path if is_hub_repo else None,
+ "source_revision": revision,
+ "configuration": self.name,
+ "split": self.split,
+ "streaming": True,
+ }
+ if state is not None:
+ state.set(state_key, result)
+ self._emit_asset_event(context, result)
+ return result
+
+ def _emit_asset_event(self, context: Context, publication: dict[str, Any])
-> None:
+ events = context.get("outlet_events")
+ if events is not None:
+ events[self.published_asset].extra = {
+ "repository_id": self.repo_id,
+ "commit_sha": publication["commit_sha"],
+ "commit_url": publication["commit_url"],
+ "source_revision": publication["source_revision"],
+ }
+
+
+class MapDatasetOperator(BaseOperator):
+ """Apply a Python callable and persist the transformed artifact."""
+
+ template_fields = ("input_uri", "output_uri", "transform_kwargs",
"cache_dir")
+
+ def __init__(
+ self,
+ *,
+ input_uri: str | dict[str, Any],
+ output_uri: str,
+ transform: Callable[..., Any],
+ transform_kwargs: dict[str, Any] | None = None,
+ map_kwargs: dict[str, Any] | None = None,
+ cache_dir: str | None = None,
+ storage_conn_id: str | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ if not callable(transform):
+ raise TypeError("transform must be a callable")
+ self.input_uri, self.output_uri, self.transform = input_uri,
output_uri, transform
+ self.transform_kwargs, self.map_kwargs = transform_kwargs or {},
map_kwargs or {}
+ self.cache_dir = cache_dir
+ self.storage_conn_id = storage_conn_id
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ dataset = _load_artifact(
+ self.input_uri, cache_dir=self.cache_dir,
storage_conn_id=self.storage_conn_id
+ )
+ mapped = dataset.map(self.transform, fn_kwargs=self.transform_kwargs,
**self.map_kwargs)
+ source = self.input_uri if isinstance(self.input_uri, dict) else {}
+ provenance = {
+ key: source[key]
+ for key in ("repository_id", "revision", "configuration", "split")
+ if key in source
+ }
+ return _save(mapped, self.output_uri, provenance,
storage_conn_id=self.storage_conn_id)
+
+
+class PushToHubOperator(BaseOperator):
+ """Synchronously publish an artifact and emit its successful asset
event."""
+
+ template_fields = ("input_uri", "config", "commit_message")
+
+ def __init__(
+ self,
+ *,
+ input_uri: str | dict[str, Any],
+ repo_id: str,
+ config: str | None = None,
+ private: bool | None = None,
+ commit_message: str | None = None,
+ huggingface_conn_id: str | None = "huggingface_default",
+ storage_conn_id: str | None = None,
+ push_kwargs: dict[str, Any] | None = None,
+ durable: bool = True,
+ **kwargs: Any,
+ ) -> None:
+ if "{{" in repo_id:
+ raise ValueError("repo_id cannot be templated: it defines the
Asset that downstream DAGs schedule on.")
+ self.input_uri, self.repo_id, self.config = input_uri, repo_id, config
+ self.private, self.commit_message = private, commit_message
+ self.huggingface_conn_id, self.push_kwargs = huggingface_conn_id,
push_kwargs or {}
+ self.storage_conn_id = storage_conn_id
+ self.durable = durable
+ self.published_asset = Asset(uri=f"hf://datasets/{repo_id}")
+ user_outlets = kwargs.pop("outlets", None)
+ if user_outlets is None:
+ user_outlets = []
+ elif not isinstance(user_outlets, (list, tuple)):
+ user_outlets = [user_outlets]
+ super().__init__(outlets=[*user_outlets, self.published_asset],
**kwargs)
Review Comment:
When this is mapped with `.partial().expand()`, `__init__` doesn't run at
parse time and `MappedOperator.outlets` reads only `partial_kwargs`, so the
`hf://datasets/...` asset never reaches the serialized Dag and the event
emitted at runtime has nothing to attach to. `StreamDatasetOperator` has the
same shape. Is mapping meant to work here?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Sequence,
Value
+
+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
+
+
+def _sha256_path(uri: str, *, conn_id: str | None = None) -> str:
+ digest = hashlib.sha256()
+ with _path(uri, conn_id=conn_id).open("rb") as stream:
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _validate_feature(feature: Any, field_name: str) -> None:
+ """Allow Parquet-safe scalar, class-label, and nested sequence features."""
+ if isinstance(feature, (Value, ClassLabel)):
+ return
+ if isinstance(feature, Sequence):
Review Comment:
`LargeList` isn't a `Sequence` subclass, so a dataset with `large_list`
columns (Polars writes those) fails here, after the whole download. Can
`LargeList` be accepted next to `Sequence`?
##########
providers/huggingface/tests/unit/huggingface/operators/test_datasets_operators.py:
##########
@@ -0,0 +1,546 @@
+# 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 datetime import datetime, timezone
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Image, Value
+from huggingface_hub import CommitInfo, GitCommitInfo
+
+from airflow.providers.huggingface.operators.datasets import (
+ LoadDatasetOperator,
+ MapDatasetOperator,
+ PushToHubOperator,
+ StreamDatasetOperator,
+ _is_hub_repo,
+ _publication_marker,
+ _load_artifact,
+ _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=timezone.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("airflow.providers.huggingface.hooks.huggingface.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("airflow.providers.huggingface.hooks.huggingface.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,
+ )
+ result = operator.execute({"run_id": "scheduled__2026-10-31"})
+
+ 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_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_called_once_with(
+ "org/data", repo_type="dataset", revision="main"
+ )
+ 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) 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, _ = 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}
+ )
+ 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]
+ 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
+ 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}
+ )
+ 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"
+ )
+ with
patch("airflow.providers.huggingface.operators.datasets._load_artifact") as
load:
+ result = operator.execute({})
+ 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}
+ )
+ 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.side_effect = [[], [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")
+ class FailingStore:
+ def get(self, key): return None
+ def set(self, key, value): 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 None
+ def set(self, key, value): pass
+ 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}
+ )
+ 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 1-of-2)")
+ ]
+ 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"
+ )
+ result = operator.execute({})
+
+ 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)]
+ 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 = None
+ hook.push_to_hub.side_effect = commits
+ configs = [
+ ("config-a", {"revision": "branch-a"}),
+ ("config-b", {"revision": "branch-a"}),
+ ("config-b", {"revision": "branch-b", "create_pr": True}),
+ ]
+ for index, (config, push_kwargs) in enumerate(configs):
+ operator = PushToHubOperator(
+ task_id=f"publish_{index}",
+ input_uri=reference,
+ repo_id="org/data",
+ config=config,
+ push_kwargs=push_kwargs,
+ huggingface_conn_id=None,
+ )
+ operator.execute({})
+
+ markers = [call.args[1] for call in
hook.find_commit_by_marker.call_args_list]
+ assert len(set(markers)) == 3
+ assert
hook.find_commit_by_marker.call_args_list[2].kwargs["include_pull_requests"] is
True
+
+
+def test_load_retry_reuses_resolved_revision_from_task_state_store(tmp_path,
monkeypatch):
+ monkeypatch.setenv("AIRFLOW_CONN_HUGGINGFACE_TEST",
'{"conn_type":"huggingface","password":"operator-test-token"}')
+ dataset = Dataset.from_dict({"value": [1]})
+ store_data = {}
+ class Store:
+ def get(self, key): return store_data.get(key)
+ def set(self, key, value): store_data[key] = value
+ ti = type("TaskInstance", (), {"dag_id": "d", "task_id": "load"})()
+ context = {"task_state_store": Store(), "ti": ti}
+ with (
+ patch("airflow.providers.huggingface.hooks.huggingface.HfApi",
autospec=True) as api,
+ patch("airflow.providers.huggingface.hooks.huggingface.load_dataset",
autospec=True) as load_dataset,
+ ):
+ api.return_value.dataset_info.return_value = type("Info", (), {"sha":
"pinned-sha"})()
+ load_dataset.return_value = dataset
+ operator = LoadDatasetOperator(
+ task_id="load", path="org/data", revision="main",
+ output_uri=(tmp_path / "first.json").as_uri(),
+ huggingface_conn_id="huggingface_test",
+ )
+ first_result = operator.execute(context)
+ operator.output_uri = (tmp_path / "retry.json").as_uri()
+ retry_result = operator.execute(context)
+ assert first_result["revision"] == retry_result["revision"] == "pinned-sha"
+ api.return_value.dataset_info.assert_called_once_with("org/data",
revision="main")
+ assert [call.kwargs["revision"] for call in load_dataset.call_args_list]
== [
+ "pinned-sha",
+ "pinned-sha",
+ ]
+
+
+def test_load_durable_false_does_not_read_or_write_task_state(tmp_path,
monkeypatch):
+ monkeypatch.setenv("AIRFLOW_CONN_HUGGINGFACE_TEST",
'{"conn_type":"huggingface","password":"operator-test-token"}')
+ with (
+ patch("airflow.providers.huggingface.hooks.huggingface.HfApi",
autospec=True) as api,
+ patch("airflow.providers.huggingface.hooks.huggingface.load_dataset",
autospec=True) as load_dataset,
+ ):
+ api.return_value.dataset_info.side_effect = [
+ type("Info", (), {"sha": "first-sha"})(),
+ type("Info", (), {"sha": "second-sha"})(),
+ ]
+ load_dataset.return_value = Dataset.from_dict({"value": [1]})
+ operator = LoadDatasetOperator(
+ task_id="load", path="org/data", revision="main",
+ output_uri=(tmp_path / "first.json").as_uri(),
+ huggingface_conn_id="huggingface_test", durable=False,
+ )
+ store = MagicMock()
+ context = {"task_state_store": store}
+ first = operator.execute(context)
+ operator.output_uri = (tmp_path / "second.json").as_uri()
+ second = operator.execute(context)
+ assert first["revision"] == "first-sha"
+ assert second["revision"] == "second-sha"
+ store.get.assert_not_called()
+ store.set.assert_not_called()
+ assert api.return_value.dataset_info.call_count == 2
+
+
+def
test_push_preserves_user_outlets_and_changes_commit_for_new_fingerprint(tmp_path,
hub_dependencies):
+ hf_api, push_to_hub = hub_dependencies
+ reference = _save(Dataset.from_dict({"value": [1]}), (tmp_path /
"input.json").as_uri())
+ caller_asset = type("Asset", (), {"uri": "custom://caller"})()
+ first_commit = commit_info(oid="1" * 40)
+ second_commit = commit_info(oid="2" * 40)
+ push_to_hub.side_effect = [first_commit, second_commit]
+ state_data = {}
+
+ class Store:
+ def get(self, key):
+ return state_data.get(key)
+
+ def set(self, key, value):
+ state_data[key] = value
+
+ operator = PushToHubOperator(task_id="publish", input_uri=reference,
repo_id="org/data",
+ huggingface_conn_id="huggingface_test",
outlets=[caller_asset])
+ assert caller_asset in operator.outlets
+ assert operator.published_asset in operator.outlets
+ operator.execute({"task_state_store": Store()})
+ changed = _save(Dataset.from_dict({"value": [2]}), (tmp_path /
"changed.json").as_uri())
+ operator.input_uri = changed
+ result = operator.execute({"task_state_store": Store()})
+ assert push_to_hub.call_count == 2
+ assert result["commit_sha"] == second_commit.oid
+
+
+def test_push_rejects_templated_repo_id():
+ with pytest.raises(ValueError, match="repo_id cannot be templated"):
+ PushToHubOperator(task_id="publish", input_uri={}, repo_id="org/{{
params.repo }}")
+
+
+def test_existing_local_path_is_not_treated_as_hub_repo(tmp_path):
+ local_path = tmp_path / "datasets" / "local"
+ local_path.mkdir(parents=True)
+ assert _is_hub_repo(str(local_path)) is False
+
+ with (
+
patch("airflow.providers.huggingface.operators.datasets.HuggingFaceHook",
autospec=True) as hook_type,
+ patch("airflow.providers.huggingface.operators.datasets._save") as
save,
+ ):
+ hook = hook_type.return_value
+ hook.load_dataset.return_value = Dataset.from_dict({"value": [1]})
+ save.return_value = {"revision": None, "repository_id": None}
+ operator = LoadDatasetOperator(
+ task_id="load_local", path=str(local_path),
output_uri="file:///tmp/local.json"
+ )
+ result = operator.execute({})
+ hook.get_dataset_info.assert_not_called()
+ hook.load_dataset.assert_called_once_with(str(local_path),
download_mode="reuse_dataset_if_exists")
+ assert result["revision"] is None
+ assert result["repository_id"] is None
+
+
[email protected]
+def attached_memory_filesystem(monkeypatch):
+ from airflow.sdk.io import attach
+ from airflow.sdk.io.store import _STORE_CACHE
+ from fsspec.implementations.memory import MemoryFileSystem
+
+ monkeypatch.delitem(_STORE_CACHE, "memory", raising=False)
Review Comment:
With `raising=False`, `delitem` records nothing when `memory` isn't cached
yet (the usual case), so the store `attach` adds below stays in `_STORE_CACHE`
for later tests. Could the fixture `yield` the store and pop `memory`
afterwards?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Sequence,
Value
+
+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
+
+
+def _sha256_path(uri: str, *, conn_id: str | None = None) -> str:
+ digest = hashlib.sha256()
+ with _path(uri, conn_id=conn_id).open("rb") as stream:
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _validate_feature(feature: Any, field_name: str) -> None:
+ """Allow Parquet-safe scalar, class-label, and nested sequence features."""
+ if isinstance(feature, (Value, ClassLabel)):
+ return
+ if isinstance(feature, Sequence):
+ _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 features"
+ )
+
+
+def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str,
Dataset]]:
+ 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."""
+ 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:
+ parquet.write_table(part.data.table, stream)
Review Comment:
If a transform or filter leaves a split with zero rows, this writes an empty
Parquet file that `Dataset.from_parquet` can't read back (`ArrowInvalid:
BatchSize must be greater than 0` on datasets 4.0.0). The writing task succeeds
and every downstream Map or Push fails, on every retry. Should `_save` reject
empty splits, or `_load_artifact` build them from the features?
##########
providers/huggingface/docs/operators.rst:
##########
@@ -0,0 +1,166 @@
+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. 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.
+
+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 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.
+
+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, it pins the resolved Hub source revision and reuses a
+completed push from task state or the matching final Hub commit on retry. 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 and accepts only a final upload commit;
intermediate
+parts of a multi-commit upload do not count as completed publication.
+
+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 task and run IDs, source and pinned revision, configuration and split,
+ destination repository, target config/revision, and pull-request setting.
+ With durability enabled, it first checks task state, then searches the
target
+ Hub refs for a completed commit carrying that marker. A match skips loading
+ the source and pushing again; the prior result and asset event are reused.
+#. **Create a lazy stream if recovery found no commit.** The operator calls
+ ``load_dataset(..., streaming=True)`` and passes the optional configuration,
+ split, and pinned revision. Hugging Face Datasets reads rows incrementally.
+ The optional ``transform`` is applied with ``IterableDataset.map`` and
+ ``transform_kwargs``; the optional ``predicate`` is then applied with
+ ``IterableDataset.filter``. Neither operation materializes the entire
+ dataset first.
+#. **Publish directly to the Hub.** The transformed iterable is passed to
+ ``push_to_hub`` in the same task. ``config``, ``private``, and
``push_kwargs``
+ control the destination. The commit message includes the recovery marker.
+ Datasets uploads Parquet shards progressively; the operator does not collect
+ the rows into a list or write an intermediate artifact.
+#. **Recover and report the commit.** If the push returns a commit, its SHA and
+ URL are returned. If it returns no commit object and durability is enabled,
+ the operator searches the Hub again in case the upload succeeded before the
+ worker stopped. It accepts only a final upload part. If no completed commit
+ can be found, the task raises ``RuntimeError``. The result contains commit
+ metadata and source provenance, not dataset rows; that result is saved to
+ task state when available and copied to the outlet event's extra fields.
+
+Unlike the local Parquet artifact path, this operator delegates feature
+serialization to Datasets' Hub upload implementation and does not use the
+provider's ``_validate_feature`` whitelist. It does not create a local manifest
+reference for downstream ``MapDatasetOperator`` or ``PushToHubOperator`` tasks;
+downstream tasks should load the published Hub repository. Passing
+``streaming=True`` through ``LoadDatasetOperator.load_kwargs`` is not a
+substitute: that operator writes a durable artifact and its writer accepts
+``Dataset`` and ``DatasetDict``, not ``IterableDataset``.
+
+The system example demonstrates the durable artifact flow and its task handoff:
Review Comment:
This is the `StreamDatasetOperator` section, but the included snippet is the
Load and Map artifact flow, which the paragraph above says this operator
doesn't use. Could this show a stream example instead, or move under the
Load/Map sections?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Sequence,
Value
+
+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
+
+
+def _sha256_path(uri: str, *, conn_id: str | None = None) -> str:
+ digest = hashlib.sha256()
+ with _path(uri, conn_id=conn_id).open("rb") as stream:
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _validate_feature(feature: Any, field_name: str) -> None:
+ """Allow Parquet-safe scalar, class-label, and nested sequence features."""
+ if isinstance(feature, (Value, ClassLabel)):
+ return
+ if isinstance(feature, Sequence):
+ _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 features"
+ )
+
+
+def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str,
Dataset]]:
+ 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."""
+ 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:
+ parquet.write_table(part.data.table, 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": _sha256_path(split_uri, conn_id=storage_conn_id),
+ }
+
+ 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:
+ try:
+ path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ try:
+ manifest_path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ 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:
+ 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})
+ raw_split = item.get("split") or (name if name != "dataset" else
None)
+ split_name = raw_split.split("[")[0] if raw_split and "[" in
raw_split else raw_split
+ restored = Dataset.from_parquet(
+ str(local_path),
+ features=features,
+ split=split_name,
+ 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=(",",
":")).encode()
+ return f"{prefix}:{hashlib.sha256(payload).hexdigest()}"
+
+
+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:
+ """Match a Hub dataset repo ID without treating an existing local path as
remote."""
+ return path.count("/") == 1 and not path.startswith((".", "/")) and not
Path(path).exists()
+
+
+class LoadDatasetOperator(BaseOperator):
+ """Load a dataset revision and persist it as a Parquet artifact."""
+
+ template_fields = ("path", "name", "split", "revision", "cache_dir",
"output_uri")
+
+ 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, self.output_uri, self.name, self.split = path, output_uri,
name, split
+ self.revision, self.cache_dir, self.download_mode = revision,
cache_dir, download_mode
+ self.storage_conn_id = storage_conn_id
+ self.huggingface_conn_id, self.load_kwargs = huggingface_conn_id,
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
+ state_identity = json.dumps([self.path, requested_revision],
separators=(",", ":"))
+ state_key = _state_key("load-revision", state_identity)
+ revision = state.get(state_key) if state is not None else None
+ if revision is None:
+ revision = requested_revision
+ if is_hub_repo:
+ revision = hook.get_dataset_info(self.path,
requested_revision).sha
+ if state is not None and revision is not None:
+ state.set(state_key, revision)
+ elif state is not None:
+ self.log.info("Reusing pinned Hugging Face dataset revision %s for
%s", revision, self.path)
+
+ kwargs = dict(self.load_kwargs)
+ kwargs["download_mode"] = self.download_mode
+ if self.name is not None:
+ kwargs["name"] = self.name
+ if self.split is not None:
+ kwargs["split"] = self.split
+ if revision is not None:
+ kwargs["revision"] = revision
+ if self.cache_dir is not None:
+ kwargs["cache_dir"] = self.cache_dir
+ dataset = hook.load_dataset(self.path, **kwargs)
+ metadata = {
+ "repository_id": self.path if is_hub_repo else None,
+ "revision": revision,
+ "configuration": self.name,
+ "split": self.split,
+ }
+ return _save(dataset, self.output_uri, metadata,
storage_conn_id=self.storage_conn_id)
+
+
+class StreamDatasetOperator(BaseOperator):
+ """Transform a streamed dataset lazily and publish it directly to the
Hub."""
+
+ template_fields = ("path", "name", "split", "revision", "config",
"commit_message")
+
+ def __init__(
+ self,
+ *,
+ path: str,
+ repo_id: str,
+ transform: Callable[..., dict[str, Any]] | None = None,
+ predicate: Callable[[dict[str, Any]], bool] | None = None,
+ name: str | None = None,
+ split: str | None = None,
+ revision: str | None = None,
+ config: str | None = None,
+ private: bool | None = None,
+ commit_message: str | None = None,
+ transform_kwargs: dict[str, Any] | None = None,
+ push_kwargs: dict[str, Any] | None = None,
+ huggingface_conn_id: str | None = "huggingface_default",
+ durable: bool = True,
+ **kwargs: Any,
+ ) -> None:
+ if "{{" in repo_id:
+ raise ValueError("repo_id cannot be templated: it defines the
Asset that downstream DAGs schedule on.")
+ if transform is not None and not callable(transform):
+ raise TypeError("transform must be a callable or None")
+ if predicate is not None and not callable(predicate):
+ raise TypeError("predicate must be a callable or None")
+ self.path, self.repo_id = path, repo_id
+ self.transform, self.predicate = transform, predicate
+ self.name, self.split, self.revision = name, split, revision
+ self.config, self.private, self.commit_message = config, private,
commit_message
+ self.transform_kwargs = transform_kwargs or {}
+ self.push_kwargs = push_kwargs or {}
+ self.huggingface_conn_id = huggingface_conn_id
+ self.durable = durable
+ self.published_asset = Asset(uri=f"hf://datasets/{repo_id}")
+ user_outlets = kwargs.pop("outlets", None)
+ if user_outlets is None:
+ user_outlets = []
+ elif not isinstance(user_outlets, (list, tuple)):
+ user_outlets = [user_outlets]
+ super().__init__(outlets=[*user_outlets, self.published_asset],
**kwargs)
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ hook = HuggingFaceHook(self.huggingface_conn_id)
+ is_hub_repo = _is_hub_repo(self.path)
+ requested_revision = self.revision
+ state = context.get("task_state_store") if self.durable else None
+ revision_identity = json.dumps(
+ [self.path, requested_revision, self.name, self.split],
separators=(",", ":")
+ )
+ revision_key = _state_key("stream-revision", revision_identity)
+ revision = state.get(revision_key) if state is not None else None
+ if revision is None:
+ revision = requested_revision
+ if is_hub_repo:
+ revision = hook.get_dataset_info(self.path,
requested_revision).sha
+ if state is not None and revision is not None:
+ state.set(revision_key, revision)
+ elif state is not None:
+ self.log.info("Reusing pinned Hugging Face streaming revision %s
for %s", revision, self.path)
+
+ run_id = context.get("run_id", "manual")
+ push_kwargs = dict(self.push_kwargs)
+ if self.config is not None:
+ push_kwargs["config_name"] = self.config
+ if self.private is not None:
+ push_kwargs["private"] = self.private
+ marker = _publication_marker(
+ "huggingface-stream-publish",
+ {
+ "task_id": self.task_id,
+ "run_id": run_id,
+ "source": self.path,
+ "source_revision": revision,
+ "name": self.name,
+ "split": self.split,
+ "repo_id": self.repo_id,
+ "config_name": push_kwargs.get("config_name", "default"),
+ "revision": push_kwargs.get("revision", "main"),
+ "create_pr": bool(push_kwargs.get("create_pr", False)),
+ },
+ )
+ state_key = _state_key("stream-publication", marker)
+ if state is not None:
+ prior_result = state.get(state_key)
+ if prior_result:
+ self.log.info("Reusing stored streaming publication result for
%s", self.repo_id)
+ self._emit_asset_event(context, prior_result)
+ return prior_result
+
+ target_revision = push_kwargs.get("revision")
+ include_pull_requests = bool(push_kwargs.get("create_pr", False))
+ commit = None
+ if self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=include_pull_requests,
+ )
+ if commit is not None:
+ self.log.info("Reusing existing completed streaming
publication to %s", self.repo_id)
+
+ if commit is None:
+ load_kwargs: dict[str, Any] = {"streaming": True}
+ if self.name is not None:
+ load_kwargs["name"] = self.name
+ if self.split is not None:
+ load_kwargs["split"] = self.split
+ if revision is not None:
+ load_kwargs["revision"] = revision
+ stream = hook.load_dataset(self.path, **load_kwargs)
+
+ if self.transform is not None:
+ stream = stream.map(self.transform,
fn_kwargs=self.transform_kwargs)
+ if self.predicate is not None:
+ stream = stream.filter(self.predicate)
+
+ message = self.commit_message or "Stream dataset from Airflow"
+ push_kwargs["commit_message"] = f"{message} [{marker}]"
+ commit = hook.push_to_hub(stream, self.repo_id, **push_kwargs)
+ # A retry can recover the completed Hub commit if the worker exited
+ # after upload but before Airflow recorded the task result.
+ if commit is None:
+ if self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=include_pull_requests,
+ )
+
+ if commit is None:
+ raise RuntimeError(
+ f"Hub push to {self.repo_id!r} completed without returning or
exposing a commit for marker {marker}"
+ )
+ commit_sha, commit_url = _commit_details(commit, self.repo_id, hook)
+ result = {
+ "repository_id": self.repo_id,
+ "commit_sha": commit_sha,
+ "commit_url": commit_url,
+ "url": f"https://huggingface.co/datasets/{self.repo_id}";,
+ "source_repository_id": self.path if is_hub_repo else None,
+ "source_revision": revision,
+ "configuration": self.name,
+ "split": self.split,
+ "streaming": True,
+ }
+ if state is not None:
+ state.set(state_key, result)
+ self._emit_asset_event(context, result)
+ return result
+
+ def _emit_asset_event(self, context: Context, publication: dict[str, Any])
-> None:
+ events = context.get("outlet_events")
+ if events is not None:
+ events[self.published_asset].extra = {
+ "repository_id": self.repo_id,
+ "commit_sha": publication["commit_sha"],
+ "commit_url": publication["commit_url"],
+ "source_revision": publication["source_revision"],
+ }
+
+
+class MapDatasetOperator(BaseOperator):
+ """Apply a Python callable and persist the transformed artifact."""
+
+ template_fields = ("input_uri", "output_uri", "transform_kwargs",
"cache_dir")
+
+ def __init__(
+ self,
+ *,
+ input_uri: str | dict[str, Any],
+ output_uri: str,
+ transform: Callable[..., Any],
+ transform_kwargs: dict[str, Any] | None = None,
+ map_kwargs: dict[str, Any] | None = None,
+ cache_dir: str | None = None,
+ storage_conn_id: str | None = None,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ if not callable(transform):
+ raise TypeError("transform must be a callable")
+ self.input_uri, self.output_uri, self.transform = input_uri,
output_uri, transform
+ self.transform_kwargs, self.map_kwargs = transform_kwargs or {},
map_kwargs or {}
+ self.cache_dir = cache_dir
+ self.storage_conn_id = storage_conn_id
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ dataset = _load_artifact(
+ self.input_uri, cache_dir=self.cache_dir,
storage_conn_id=self.storage_conn_id
+ )
+ mapped = dataset.map(self.transform, fn_kwargs=self.transform_kwargs,
**self.map_kwargs)
+ source = self.input_uri if isinstance(self.input_uri, dict) else {}
+ provenance = {
+ key: source[key]
+ for key in ("repository_id", "revision", "configuration", "split")
+ if key in source
+ }
+ return _save(mapped, self.output_uri, provenance,
storage_conn_id=self.storage_conn_id)
+
+
+class PushToHubOperator(BaseOperator):
+ """Synchronously publish an artifact and emit its successful asset
event."""
+
+ template_fields = ("input_uri", "config", "commit_message")
+
+ def __init__(
+ self,
+ *,
+ input_uri: str | dict[str, Any],
+ repo_id: str,
+ config: str | None = None,
+ private: bool | None = None,
+ commit_message: str | None = None,
+ huggingface_conn_id: str | None = "huggingface_default",
+ storage_conn_id: str | None = None,
+ push_kwargs: dict[str, Any] | None = None,
+ durable: bool = True,
+ **kwargs: Any,
+ ) -> None:
+ if "{{" in repo_id:
+ raise ValueError("repo_id cannot be templated: it defines the
Asset that downstream DAGs schedule on.")
+ self.input_uri, self.repo_id, self.config = input_uri, repo_id, config
+ self.private, self.commit_message = private, commit_message
+ self.huggingface_conn_id, self.push_kwargs = huggingface_conn_id,
push_kwargs or {}
+ self.storage_conn_id = storage_conn_id
+ self.durable = durable
+ self.published_asset = Asset(uri=f"hf://datasets/{repo_id}")
+ user_outlets = kwargs.pop("outlets", None)
+ if user_outlets is None:
+ user_outlets = []
+ elif not isinstance(user_outlets, (list, tuple)):
+ user_outlets = [user_outlets]
+ super().__init__(outlets=[*user_outlets, self.published_asset],
**kwargs)
+
+ def execute(self, context: Context) -> dict[str, Any]:
+ source = self.input_uri if isinstance(self.input_uri, dict) else {}
+ fingerprint = source.get("fingerprint")
+ if not fingerprint:
+ raise ValueError("Publication requires an artifact reference with
a fingerprint")
+ kwargs = dict(self.push_kwargs)
+ if self.config is not None:
+ kwargs["config_name"] = self.config
+ if self.private is not None:
+ kwargs["private"] = self.private
+ target_revision = kwargs.get("revision")
+ create_pr = bool(kwargs.get("create_pr", False))
+ target_marker = _publication_marker(
+ "target",
+ {
+ "repo_id": self.repo_id,
+ "config_name": kwargs.get("config_name", "default"),
+ "revision": target_revision or "main",
+ "create_pr": create_pr,
+ },
+ )
+ marker = f"airflow-artifact-sha256:{fingerprint}:{target_marker}"
+ state = context.get("task_state_store") if self.durable else None
+ state_key = _state_key("publication", marker)
+ if state is not None:
+ prior_result = state.get(state_key)
+ if prior_result:
+ self.log.info("Reusing stored publication result for %s",
self.repo_id)
+ self._emit_asset_event(context, prior_result)
+ return prior_result
+
+ hook = HuggingFaceHook(self.huggingface_conn_id)
+ commit = None
+ if self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=create_pr,
+ )
+ if commit is not None:
+ self.log.info("Reusing existing completed publication to %s",
self.repo_id)
+ if commit is None:
+ dataset = _load_artifact(self.input_uri,
storage_conn_id=self.storage_conn_id)
+ message = self.commit_message or "Publish dataset from Airflow"
+ kwargs["commit_message"] = f"{message} [{marker}]"
+ commit = hook.push_to_hub(dataset, self.repo_id, **kwargs)
+ if commit is None and self.durable:
+ commit = hook.find_commit_by_marker(
+ self.repo_id,
+ marker,
+ revision=target_revision,
+ include_pull_requests=create_pr,
+ )
+
+ if commit is None:
+ raise RuntimeError(
+ f"Hub push to {self.repo_id!r} completed without returning or
exposing a final commit for {marker}"
+ )
+ commit_sha, commit_url = _commit_details(commit, self.repo_id, hook)
+ result = {
+ "repository_id": self.repo_id,
+ "commit_sha": commit_sha,
+ "commit_url": commit_url,
+ "url": f"https://huggingface.co/datasets/{self.repo_id}";,
+ "fingerprint": fingerprint,
+ "source_repository_id": source.get("repository_id"),
+ "source_revision": source.get("revision"),
+ }
+ if state is not None:
+ state.set(state_key, result)
+ self._emit_asset_event(context, result)
+ return result
+
+ def _emit_asset_event(self, context: Context, publication: dict[str, Any])
-> None:
+ events = context.get("outlet_events")
+ if events is not None:
+ events[self.published_asset].extra = {
Review Comment:
With `create_pr=True` or a non-main `revision`, this still emits the
repo-level asset event, and the extra has no ref or PR URL. Downstream Dags
scheduled on `hf://datasets/<repo>` then load `main`, which didn't change.
Should PR pushes skip the event, or should the extra carry the revision?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
+from datasets import ClassLabel, Dataset, DatasetDict, Features, Sequence,
Value
+
+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
+
+
+def _sha256_path(uri: str, *, conn_id: str | None = None) -> str:
+ digest = hashlib.sha256()
+ with _path(uri, conn_id=conn_id).open("rb") as stream:
+ for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _validate_feature(feature: Any, field_name: str) -> None:
+ """Allow Parquet-safe scalar, class-label, and nested sequence features."""
+ if isinstance(feature, (Value, ClassLabel)):
+ return
+ if isinstance(feature, Sequence):
+ _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 features"
+ )
+
+
+def _dataset_parts(dataset: Dataset | DatasetDict) -> list[tuple[str,
Dataset]]:
+ 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."""
+ 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:
+ parquet.write_table(part.data.table, 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": _sha256_path(split_uri, conn_id=storage_conn_id),
+ }
+
+ 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:
+ try:
+ path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ try:
+ manifest_path.unlink(missing_ok=True)
+ except (OSError, NotImplementedError):
+ pass
+ 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:
+ 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})
+ raw_split = item.get("split") or (name if name != "dataset" else
None)
+ split_name = raw_split.split("[")[0] if raw_split and "[" in
raw_split else raw_split
+ restored = Dataset.from_parquet(
+ str(local_path),
+ features=features,
+ split=split_name,
+ 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=(",",
":")).encode()
+ return f"{prefix}:{hashlib.sha256(payload).hexdigest()}"
+
+
+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:
+ """Match a Hub dataset repo ID without treating an existing local path as
remote."""
+ return path.count("/") == 1 and not path.startswith((".", "/")) and not
Path(path).exists()
Review Comment:
Hub IDs without a namespace, like `imdb` or `squad`, have no `/`, so this
returns False even though `load_dataset` resolves them from the Hub. Those
loads skip the revision pin and record `repository_id: None`. Could anything
that isn't an existing local path or a packaged builder name (`csv`, `parquet`,
...) count as a Hub ID?
##########
providers/huggingface/src/airflow/providers/huggingface/operators/datasets.py:
##########
@@ -0,0 +1,639 @@
+# 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 json
+from pathlib import Path
+import tempfile
+from typing import Any, Callable
+
+import pyarrow.parquet as parquet
Review Comment:
Importing `datasets` and `pyarrow` at module level takes about 2.7 s and 125
MB per process on datasets 4.0.0, and every Dag file that imports these
operators pays that at parse time (the hook module imports `datasets` too).
Could they move into the methods that use them, with `TYPE_CHECKING` imports
for the annotations?
--
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]
