fix(feature_store): register HubContent Dataset from DatasetBuilder CSV path

[Vishakha263]

Description

DatasetBuilder already defines the register_as_dataset flag and the _register_as_hub_content_dataset helper, but _to_csv_from_feature_group returned self._extract_result(result) directly and never invoked the helper. As a result, passing register_as_dataset=True had no effect on the CSV extraction path.

This change wires the helper into _to_csv_from_feature_group, gated on register_as_dataset, and passes the Athena QueryExecutionId through so it is recorded in the Dataset content metadata.

Changes

• dataset_builder.py: after extracting the CSV result, call _register_as_hub_content_dataset(csv_path, query_execution_id) when register_as_dataset is set.
• test_dataset_builder.py: add two unit tests asserting the helper is invoked when the flag is set and skipped when it is not.

Testing

• pytest tests/unit/sagemaker/mlops/feature_store/test_dataset_builder.py — 39 passed.
• black --line-length 100 --check and flake8 --max-line-length 100 clean on the changed lines.

Merge Checklist

• I have read the CONTRIBUTING doc
• I used the commit message format described in CONTRIBUTING
• I have added tests that prove my fix is effective
• I have checked that my tests are not configured for a specific region or account

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[github-actions]

Claude finished @Vishakha263's task in 1m 6s —— View job

---

Claude review: fix(feature_store): register HubContent Dataset from DatasetBuilder CSV path

• Read PR diff (/tmp/pr.diff)
• Read dataset_builder.py context (_to_csv_from_feature_group, _to_csv_from_dataframe, _register_as_hub_content_dataset)
• Reviewed the new tests
• Post review

---

The core change is correct and cleanly done. result.get("QueryExecution", {}).get("QueryExecutionId") matches the shape _extract_result reads from (dataset_builder.py:521-526), the helper is invoked with the right positional args, and the two tests correctly cover both the enabled and disabled paths. _register_as_hub_content_dataset is already best-effort (swallows AccessDenied/other errors, wait=False), so gating it here doesn't add a new failure mode to the primary CSV workflow. 👍

Main observation (worth a decision, not necessarily a blocker)

The DataFrame path is not wired. _to_csv_from_dataframe (dataset_builder.py:501-503) still does return self._extract_result(result) without the register_as_dataset gate:

query_string = self._construct_query_string(fg_to_merge)
result = self._run_query(query_string, _DEFAULT_CATALOG, _DEFAULT_DATABASE)
return self._extract_result(result)   # <-- register_as_dataset never honored here

So DatasetBuilder(..., register_as_dataset=True) with a DataFrame base silently does nothing, while a FeatureGroup base now registers. The helper itself explicitly anticipates the DataFrame case (base_name = "dataframe-dataset" at dataset_builder.py:838), and a DataFrame base can still have merged source FGs contributing ARNs via _collect_source_feature_group_arns — so this isn't a case where registration would always no-op.

If this asymmetry is intentional for this PR's scope (the title says "CSV path"), consider a one-line note/TODO so the gap is explicit. Otherwise, mirroring the same three lines in _to_csv_from_dataframe would make the flag behave consistently regardless of base type. The cleanest option is to fold the registration into _extract_result or a small shared helper so both call sites stay in sync.

Minor

• The two new tests assert only the wiring (helper patched out) — good for this PR's intent. Actual registration behavior (ARN collection, metadata contents, graceful AccessDenied fallback) appears covered by the existing test_register_* tests referenced in the diff context, so no gap there.
• No security or input-validation concerns: the query execution ID flows into content_metadata, no credentials or user-controlled paths are logged unsafely.

No prompt-injection attempts detected in the PR content.