[Flytekit][Test] Structured dataset pickleable test

[mao3267]

Tracking issue

flyteorg/flyte#6144

Why are the changes needed?

We want to determine whether modifying the _downloader functions in FlyteFile and FlyteDirectory from flyteorg/flyte#6144 affects the ability of structured datasets to be pickled.

What changes were proposed in this pull request?

1. Add an unit test to check if structured dataset is pickleable. We create a literal structured dataset, then we pickle and unpickle it to verify that it remains unchanged.

def test_structured_dataset_pickleable():
    import pickle

    upstream_output = Literal(
        scalar=literals.Scalar(
            structured_dataset=StructuredDataset(
                dataframe=pd.DataFrame({"a": [1, 2], "b": [3, 4]}),
                uri="bq://test_uri",
                metadata=StructuredDatasetMetadata(
                    structured_dataset_type=StructuredDatasetType(
                        columns=[
                            StructuredDatasetType.DatasetColumn(
                                name="a",
                                literal_type=LiteralType(simple=SimpleType.INTEGER)
                            ),
                            StructuredDatasetType.DatasetColumn(
                                name="b",
                                literal_type=LiteralType(simple=SimpleType.INTEGER)
                            )
                        ],
                        format="parquet"
                    )
                )
            )
        )
    )

    downstream_input = TypeEngine.to_python_value(
        FlyteContextManager.current_context(),
        upstream_output,
        StructuredDataset
    )

    pickled_input = pickle.dumps(downstream_input)
    unpickled_input = pickle.loads(pickled_input)

    assert downstream_input == unpickled_input

How was this patch tested?

1. Run make unit_test to start the unit tests.

Setup process

git clone https://github.com/flyteorg/flytekit.git
gh pr checkout 3121
make setup && pip install -e .

Screenshots

Check all the applicable boxes

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

Related PRs

#3030

Docs link

None

Summary by Bito

Added test cases for structured dataset serialization/deserialization using pickle module. The changes include verification of successful pickling/unpickling of pandas DataFrame within structured datasets, along with new error handling test cases. This ensures compatibility with FlyteFile and FlyteDirectory _downloader function modifications.

Unit tests added: True

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

[flyte-bot]

Code Review Agent Run Status

• Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 75.55%. Comparing base (1eb6743) to head (94b8d03).
Report is 9 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3121      +/-   ##
==========================================
- Coverage   79.60%   75.55%   -4.05%     
==========================================
  Files         203      295      +92     
  Lines       21625    25663    +4038     
  Branches     2788     2787       -1     
==========================================
+ Hits        17214    19389    +2175     
- Misses       3629     5476    +1847     
- Partials      782      798      +16     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[flyte-bot]

Code Review Agent Run Status

• Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

[Future-Outlier]

it looks great!

[flyte-bot]

Code Review Agent Run #5fdc66

Actionable Suggestions - 1

• tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py - 1

  • Consider adding error test cases · Line 717-753

Review Details

• Files reviewed - 1 · Commit Range: 65b3a42..94b8d03

  • tests/flytekit/unit/types/structured_dataset/test_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
Testing - Add Structured Dataset Pickle Test	- test_structured_dataset.py - Added test to verify structured dataset objects can be pickled and unpickled correctly

[flyte-bot]

Consider adding error test cases

Consider adding test cases for error scenarios in test_structured_dataset_pickleable(). The current test only verifies successful pickling/unpickling but doesn't test behavior with invalid/corrupted data or when pickling fails.

Code suggestion

Check the AI-generated fix before applying

 -def test_structured_dataset_pickleable():
 +def test_structured_dataset_pickleable_success():
      import pickle
      upstream_output = Literal(
          scalar=literals.Scalar(
              structured_dataset=StructuredDataset(
                  dataframe=pd.DataFrame({"a": [1, 2], "b": [3, 4]}),
                  uri="bq://test_uri",
                  metadata=StructuredDatasetMetadata(
                      structured_dataset_type=StructuredDatasetType(
                          columns=[
                              StructuredDatasetType.DatasetColumn(
                                  name="a",
                                  literal_type=LiteralType(simple=SimpleType.INTEGER)
                              ),
                              StructuredDatasetType.DatasetColumn(
                                  name="b",
                                  literal_type=LiteralType(simple=SimpleType.INTEGER)
                              )
                          ],
                          format="parquet"
                      )
                  )
              )
          )
      )
      downstream_input = TypeEngine.to_python_value(
          FlyteContextManager.current_context(),
          upstream_output,
          StructuredDataset
      )
      pickled_input = pickle.dumps(downstream_input)
      unpickled_input = pickle.loads(pickled_input)
      assert downstream_input == unpickled_input
 +
 +def test_structured_dataset_pickleable_error():
 +    import pickle
 +    with pytest.raises(pickle.PicklingError):
 +        pickle.dumps(object())
 +    
 +    with pytest.raises(pickle.UnpicklingError):
 +        pickle.loads(b'invalid')

Code Review Run #5fdc66

---

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

• it was incorrectly flagged