Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Flytekit][Test] Structured dataset pickleable test #3121

Merged
merged 2 commits into from
Feb 12, 2025
Merged

[Flytekit][Test] Structured dataset pickleable test #3121

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
65b3a42
test: Structured dataset pickleable
mao3267 Feb 8, 2025
94b8d03
Merge branch 'master' of https://github.com/mao3267/flytekit into str…
mao3267 Feb 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from flytekit.core.workflow import workflow
from flytekit.lazy_import.lazy_module import is_imported
from flytekit.models import literals
from flytekit.models.literals import StructuredDatasetMetadata
from flytekit.models.literals import StructuredDatasetMetadata, Literal
from flytekit.models.types import LiteralType, SchemaType, SimpleType, StructuredDatasetType
from flytekit.tools.translator import get_serializable
from flytekit.types.structured.structured_dataset import (
Expand Down Expand Up @@ -713,3 +713,41 @@ def mock_resolve_remote_path(flyte_uri: str) -> typing.Optional[str]:

lit = sdte.encode(ctx, sd, df_type=pd.DataFrame, protocol="bq", format="parquet", structured_literal_type=lt)
assert lit.scalar.structured_dataset.uri == "bq://blah/blah/blah"

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
Comment on lines +717 to +753
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Loading