Return StructuredDataset which is a field in a dataclass

[arbaobao]

Tracking issue

Related to #6117

Why are the changes needed?

If we wrap the StructuredDataset in a dataclass, it will fail during the to_flyte_idl conversion.

What changes were proposed in this pull request?

Before returning Literals, we check the type of python_val._literal_sd. If it is a Python native StructuredDataset, we transform it into a Literals.StructuredDataset.

How was this patch tested?

As described in #6117, an error occurs when the extract task is executed.

@dataclass
class Data:
    f: StructuredDataset


@task
def create_data() -> Data:
    return Data(f=StructuredDataset(dataframe=pd.DataFrame({"a": [5]})))


@task
def extract(d: Data) -> StructuredDataset:
    return d.f


@workflow
def example_wf() -> None:
    d = create_data()
    f = extract(d=d)

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Summary by Bito

Fixed a bug in StructuredDataset handling within dataclasses during to_flyte_idl conversion. The PR adds proper transformation of python_val._literal_sd instances into Literals.StructuredDataset, enabling tasks to successfully return StructuredDataset objects as dataclass fields.

Unit tests added: False

Estimated effort to review (1-5, lower is better): 1

[flyte-bot]

Code Review Agent Run #63793c

Actionable Suggestions - 2

• flytekit/types/structured/structured_dataset.py - 2

  • Private member access needs encapsulation · Line 738-742
  • Consider extracting literal creation logic · Line 738-742

Review Details

• Files reviewed - 2 · Commit Range: 51f6f73..a3df842

  • flytekit/core/type_engine.py
  • flytekit/types/structured/structured_dataset.py
• Files skipped - 0
• Tools

  • Whispers (Secret Scanner) - ✔︎ Successful
  • Detect-secrets (Secret Scanner) - ✔︎ Successful
  • MyPy (Static Code Analysis) - ✔︎ Successful
  • Astral Ruff (Static Code Analysis) - ✔︎ Successful

---

AI Code Review powered by

[flyte-bot]

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Bug Fix - Fix StructuredDataset Dataclass Handling	- structured_dataset.py - Added handling for StructuredDataset fields within dataclasses during type transformation

[flyte-bot]

Private member access needs encapsulation

Accessing private member '_literal_sd'. Consider using a public interface or property to access this data.

Code suggestion

Check the AI-generated fix before applying

 -            if literal_type.structured_dataset_type is not None and self._literal_sd is not None:
 -                return self._literal_sd
 -            if literal_type.structured_dataset_type is not None and self._literal_sd is None:
 +            if literal_type.structured_dataset_type is not None and self.literal_sd is not None:
 +                return self.literal_sd
 +            if literal_type.structured_dataset_type is not None and self.literal_sd is None:

Code Review Run #63793c

---

Is this a valid issue, or was it incorrectly flagged by the Agent?

• it was incorrectly flagged

[flyte-bot]

Consider extracting literal creation logic

The code block for handling StructuredDataset passed through dataclass could be simplified by extracting the literal creation logic into a helper method. This would improve code readability and maintainability.

Code suggestion

Check the AI-generated fix before applying

Suggested change

	if isinstance(python_val._literal_sd, StructuredDataset):
	sdt = StructuredDatasetType(format=python_val._literal_sd.file_format)
	metad = literals.StructuredDatasetMetadata(structured_dataset_type=sdt)
	sd_literal = literals.StructuredDataset(uri=python_val._literal_sd.uri, metadata=metad)
	return Literal(scalar=Scalar(structured_dataset=sd_literal))
	if isinstance(python_val._literal_sd, StructuredDataset):
	return self._create_structured_dataset_literal(python_val._literal_sd.uri, python_val._literal_sd.file_format)
	def _create_structured_dataset_literal(self, uri: str, file_format: str) -> Literal:
	sdt = StructuredDatasetType(format=file_format)
	metad = literals.StructuredDatasetMetadata(structured_dataset_type=sdt)
	sd_literal = literals.StructuredDataset(uri=uri, metadata=metad)
	return Literal(scalar=Scalar(structured_dataset=sd_literal))

Code Review Run #63793c

---

Is this a valid issue, or was it incorrectly flagged by the Agent?

• it was incorrectly flagged

[Future-Outlier]

it looks correct, can you provide

1. screenshot
2. add an example to integration test to test it properlly?
   test_remote.py