feat(feature-store): Add lineage registration to DatasetBuilder

[Vishakha263]

Description
Add optional lineage registration to DatasetBuilder that creates a SM Dataset (HubContent) after extracting features from Feature Groups. This enables automatic lineage tracking between Feature Groups and Training Jobs.

Motivation
SageMaker's auto-lineage currently has a gap: there's no connection between Feature Groups and Training Jobs. When a customer extracts features via DatasetBuilder.to_csv_file() and uses the resulting CSV for training, that relationship is lost. This PR fills that gap by registering the extraction as a SM Dataset with source Feature Group metadata.

Implementation details:

• Added register_as_dataset parameter to DatasetBuilder.create() (default False — fully backward compatible)
• After Athena query completes, calls HubContent.import_hub_content() with source FG ARNs in the document metadata
• Tracks source Feature Groups from both base and merged (via with_feature_group()) calls
• Deduplicates ARNs if the same FG is used multiple times
• Best-effort: if registration fails (missing permissions, service error), logs a warning and returns the CSV normally

Files changed:

• dataset_builder.py
• test_dataset_builder.py

Testing
Added 10 unit tests in TestDatasetBuilderRegisterAsDataset:

• Default flag behavior (False by default)
• Opt-in flag stored correctly
• ARN collection from base FG
• ARN collection from base + merged FGs
• ARN deduplication
• DataFrame base returns empty ARN list
• HubContent API called with correct parameters
• Graceful handling of AccessDeniedException
• Graceful handling of generic errors
• Registration skipped when no FG ARNs available
• All tests passing locally.

Permissions required (for customers who opt in)
sagemaker:ImportHubContent on the execution role. If missing, a warning is logged and the CSV is still returned.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[huiqunli]

Instead of calling HubContent.import_hub_content directly, might consider to use DataSet.create from sagemaker.ai_registry.dataset. It handles document construction and schema compliance internally. We'd need to extend DataSet.create to accept a metadata param for passing sourceFeatureGroups though.

[github-actions]

Claude finished @rsareddy0329's task in 2m 14s —— View job

---

PR Review: Add lineage registration to DatasetBuilder

• Read PR diff
• Examine dataset_builder.py context (call site, dataclass fields)
• Examine dataset.py / dataset_utils.py changes
• Review tests
• Post review feedback

Note: the inline-comment tool wasn't available in this run, so I've consolidated findings below with file/line references.

---

🔴 Blocking: the feature never actually runs

_register_as_hub_content_dataset(...) is defined (dataset_builder.py:806) and _register_as_dataset is stored (:264, :307), but nothing ever calls the registration method. A grep across sagemaker-mlops/src for _register_as_hub_content_dataset returns only its definition:

dataset_builder.py:806:    def _register_as_hub_content_dataset(

to_csv_file() (:435) returns immediately from _to_csv_from_dataframe / _to_csv_from_feature_group and never checks self._register_as_dataset or invokes registration. As written, setting register_as_dataset=True has no effect — the entire lineage feature is dead code.

The wiring needs to be added, e.g. in to_csv_file() after the CSV is produced:

csv_path, query_string = ...  # result of the two branches
if self._register_as_dataset:
    self._register_as_hub_content_dataset(csv_path)
return csv_path, query_string

All 10 new unit tests call _register_as_hub_content_dataset / _collect_source_feature_group_arns directly, so they pass without the method ever being reached through the public API. There is no test that exercises to_csv_file() with register_as_dataset=True and asserts registration happened — which is exactly why the missing call site slipped through. Please add an end-to-end test at the to_csv_file() level.

---

🟠 query_execution_id is never populated

_register_as_hub_content_dataset accepts query_execution_id and stores it in AthenaQueryExecutionId, but:

• _extract_result (:515) returns (OutputLocation, Query) — the query string, not an execution ID. The Athena QueryExecutionId is available in the query result dict but isn't threaded through.
• Since the method is never called, it always defaults to None → "".

If provenance tracking is a goal (per the PR description), plumb the actual QueryExecutionId from the Athena response through to this method. Otherwise the metadata field will always be empty.

---

🟡 Dataclass field ordering / init=False inconsistency

dataset_builder.py:264-265:

_register_as_dataset: bool = False
_source_feature_groups: List = field(default_factory=list)

Every other private, non-constructor field in this dataclass uses field(..., init=False) (:255-263). These two are left as init=True, which means they're now part of the public DatasetBuilder(...) constructor signature. _source_feature_groups in particular is internal mutable state that a caller should never set positionally. Recommend field(default=False, init=False) and field(default_factory=list, init=False) for consistency and to keep them out of the public init contract. create() sets _register_as_dataset via keyword, so it can stay a non-init field set after construction if needed, but the current split is inconsistent.

Also worth a quick check: _source_feature_groups: List has no element type — List[FeatureGroup] would match the typing style used elsewhere in the file.

---

🟡 Exception-string matching is fragile (minor, and both branches are identical)

:859 distinguishes AccessDenied via "AccessDenied" in error_msg or "not authorized" in error_msg.lower(). String-sniffing exceptions is brittle — if DataSet.create raises a botocore ClientError, inspecting e.response["Error"]["Code"] == "AccessDeniedException" is more reliable. That said, both branches only log a warning, so the distinction is purely cosmetic in the message. Either catch ClientError specifically, or simplify to a single warning. Note also that a blanket except Exception will also swallow programming errors (e.g. AttributeError, ImportError from the local import at :837) — consider narrowing.

---

🟢 Looks good

• dataset.py / dataset_utils.py changes are clean: content_metadata is optional, ContentMetadata is only added to the document when truthy (dataset_utils.py:98-99), and skipping file/format validation for metadata-only datasets is reasonable. Tests cover both the with/without-metadata document serialization and the validation-skip path.
• ARN dedup logic in _collect_source_feature_group_arns (:791) is correct, and the getattr(..., None) guards avoid AttributeError on DataFrame bases.
• Backward compatibility: register_as_dataset defaults to False, and content_metadata defaults to None, so existing callers are unaffected. ✅
• The reviewer suggestion (huiqunli) to route through DataSet.create instead of calling HubContent.import_hub_content directly was adopted. 👍

---

Summary

The DataSet.create extension and its tests are solid, but the headline feature is not wired up — register_as_dataset=True currently does nothing because _register_as_hub_content_dataset is never invoked from to_csv_file(). This is the must-fix item, along with an end-to-end test through the public API to catch it. The init=False inconsistency and query_execution_id plumbing are secondary.