[wip] Pandas Dataframe in Dataclass #3116

Future-Outlier · 2025-02-07T03:03:54Z

Tracking issue

Why are the changes needed?

we want to improve UX.

import flytekit
from flytekit import task, workflow
from dataclasses import dataclass, field
import pandas as pd
# from pydantic import BaseModel
from flytekit.types.structured.structured_dataset import (
    PARQUET,
    StructuredDataset,
    StructuredDatasetDecoder,
    StructuredDatasetEncoder,
    StructuredDatasetTransformerEngine,
)

@dataclass
class smallDC:
    df: pd.DataFrame = field(default_factory=lambda: pd.DataFrame([[1, 2, 3], [4, 5, 6]]))

@dataclass
class DC:
    dc: smallDC
    df: pd.DataFrame = field(default_factory=lambda: pd.DataFrame([[1, 2, 3], [4, 5, 6]]))

@task
def t1() -> DC:
    return DC(dc=smallDC())

if __name__ == "__main__":
    print(t1())

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

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

Related PRs

Docs link

Summary by Bito

This PR implements automatic Pandas DataFrame to StructuredDataset conversion within dataclass structures, with recursive transformation capabilities and encoding/decoding support. The DataclassTransformer in flytekit's type engine is enhanced to handle nested dataclass conversions and improved Pandas DataFrame handling through a convert_dataclass helper function. The update includes improvements to JSON decoding in binary_to_python method while removing debug-related print statements and improving documentation for dataclass case handling.

Unit tests added: False

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

Signed-off-by: Future-Outlier <[email protected]>

flyte-bot · 2025-02-07T03:04:08Z

Code Review Agent Run #fedbf7

Actionable Suggestions - 2

flytekit/core/type_engine.py - 2
- Consider dedicated transformer for DataFrame conversion · Line 882-884
- Consider caching decoder with original type · Line 1028-1029

Review Details

Files reviewed - 2 · Commit Range: 28db576..28db576
- 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 · 2025-02-07T03:07:17Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Enhanced Dataclass Support for Pandas DataFrames	- `type_engine.py` - Added support for automatic conversion of Pandas DataFrames within dataclasses to StructuredDataset format

flyte-bot · 2025-02-07T03:07:18Z

flytekit/core/type_engine.py

+            import pandas as pd
+            if isinstance(python_val, pd.DataFrame):
+                python_val = StructuredDataset(dataframe=python_val, file_format="parquet")


Consider dedicated transformer for DataFrame conversion

Consider moving the pandas DataFrame conversion logic to a dedicated transformer class instead of handling it in the _make_dataclass_serializable method. This would improve code organization and maintainability.

Code suggestion

Check the AI-generated fix before applying

Suggested change

import pandas as pd

if isinstance(python_val, pd.DataFrame):

python_val = StructuredDataset(dataframe=python_val, file_format="parquet")

Code Review Run #fedbf7

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

it was incorrectly flagged

flyte-bot · 2025-02-07T03:07:19Z

flytekit/core/type_engine.py

+                decoder = JSONDecoder(new_expected_python_type)
+                self._json_decoder[new_expected_python_type] = decoder


Consider caching decoder with original type

Consider storing the decoder in the original expected_python_type key instead of new_expected_python_type to avoid potential memory leaks from storing multiple decoders for the same logical type.

Code suggestion

Check the AI-generated fix before applying

Suggested change

decoder = JSONDecoder(new_expected_python_type)

self._json_decoder[new_expected_python_type] = decoder

decoder = JSONDecoder(new_expected_python_type)

self._json_decoder[expected_python_type] = decoder

Code Review Run #fedbf7

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

it was incorrectly flagged

Signed-off-by: Future-Outlier <[email protected]>

flyte-bot · 2025-02-07T05:18:00Z

Code Review Agent Run #6cea0d

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: 28db576..fc506fe
- 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

Signed-off-by: Future-Outlier <[email protected]>

flyte-bot · 2025-02-07T08:41:41Z

Code Review Agent Run #b16268

Actionable Suggestions - 2

flytekit/core/type_engine.py - 2
- Consider adding type hints to function · Line 972-972
- Consider validating dataclass field compatibility · Line 1032-1032

Review Details

Files reviewed - 1 · Commit Range: fc506fe..de150b1
- flytekit/core/type_engine.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 · 2025-02-07T08:44:44Z

flytekit/core/type_engine.py

+        from flytekit.types.structured.structured_dataset import StructuredDataset
+        from typing import get_type_hints, Type, Dict
+
+        def convert_dataclass(instance, target_cls):


Consider adding type hints to function

The convert_dataclass function could benefit from type hints for better code maintainability and IDE support. Consider adding type annotations for instance and target_cls parameters.

Code suggestion

Check the AI-generated fix before applying

Suggested change

def convert_dataclass(instance, target_cls):

def convert_dataclass(instance: Any, target_cls: Type[T]) -> T:

Code Review Run #b16268

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

it was incorrectly flagged

flyte-bot · 2025-02-07T08:44:45Z

flytekit/core/type_engine.py


            dc = decoder.decode(json_str)

-        return self._fix_dataclass_int(expected_python_type, dc)
+        return convert_dataclass(self._fix_dataclass_int(new_expected_python_type, dc), expected_python_type)


Consider validating dataclass field compatibility

The conversion from new_expected_python_type to expected_python_type may lose data if the types have different field structures. Consider validating field compatibility before conversion.

Code suggestion

Check the AI-generated fix before applying

Suggested change

return convert_dataclass(self._fix_dataclass_int(new_expected_python_type, dc), expected_python_type)

fixed_dc = self._fix_dataclass_int(new_expected_python_type, dc)

# Validate field compatibility

if not all(f.name in [ef.name for ef in fields(expected_python_type)] for f in fields(fixed_dc.__class__)):

raise ValueError(f"Incompatible field structure between {new_expected_python_type.__name__} and {expected_python_type.__name__}")

return convert_dataclass(fixed_dc, expected_python_type)

Code Review Run #b16268

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

it was incorrectly flagged

Pandas Dataframe in Dataclass

28db576

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla and thomasjpfan as code owners February 7, 2025 03:03

flyte-bot reviewed Feb 7, 2025

View reviewed changes

update

fc506fe

Signed-off-by: Future-Outlier <[email protected]>

add convert_dataclass in der

de150b1

Signed-off-by: Future-Outlier <[email protected]>

flyte-bot reviewed Feb 7, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[wip] Pandas Dataframe in Dataclass #3116

[wip] Pandas Dataframe in Dataclass #3116

Future-Outlier commented Feb 7, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Feb 7, 2025 •

edited

Loading

Code Review Agent Run #fedbf7

flyte-bot commented Feb 7, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Feb 7, 2025

flyte-bot Feb 7, 2025

flyte-bot commented Feb 7, 2025 •

edited

Loading

Code Review Agent Run #6cea0d

flyte-bot commented Feb 7, 2025 •

edited

Loading

Code Review Agent Run #b16268

flyte-bot Feb 7, 2025

flyte-bot Feb 7, 2025

	import pandas as pd
	if isinstance(python_val, pd.DataFrame):
	python_val = StructuredDataset(dataframe=python_val, file_format="parquet")

		decoder = JSONDecoder(new_expected_python_type)
		self._json_decoder[new_expected_python_type] = decoder

	def convert_dataclass(instance, target_cls):
	def convert_dataclass(instance: Any, target_cls: Type[T]) -> T:

-        return convert_dataclass(self._fix_dataclass_int(new_expected_python_type, dc), expected_python_type)
+        fixed_dc = self._fix_dataclass_int(new_expected_python_type, dc)
+        # Validate field compatibility
+        if not all(f.name in [ef.name for ef in fields(expected_python_type)] for f in fields(fixed_dc.__class__)):
+            raise ValueError(f"Incompatible field structure between {new_expected_python_type.__name__} and {expected_python_type.__name__}")
+        return convert_dataclass(fixed_dc, expected_python_type)

[wip] Pandas Dataframe in Dataclass #3116

Are you sure you want to change the base?

[wip] Pandas Dataframe in Dataclass #3116

Conversation

Future-Outlier commented Feb 7, 2025 • edited by flyte-bot Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

flyte-bot commented Feb 7, 2025 • edited Loading

Code Review Agent Run #fedbf7

flyte-bot commented Feb 7, 2025 • edited Loading

Changelist by Bito

flyte-bot Feb 7, 2025

Choose a reason for hiding this comment

flyte-bot Feb 7, 2025

Choose a reason for hiding this comment

flyte-bot commented Feb 7, 2025 • edited Loading

Code Review Agent Run #6cea0d

flyte-bot commented Feb 7, 2025 • edited Loading

Code Review Agent Run #b16268

flyte-bot Feb 7, 2025

Choose a reason for hiding this comment

flyte-bot Feb 7, 2025

Choose a reason for hiding this comment

Future-Outlier commented Feb 7, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Feb 7, 2025 •

edited

Loading

flyte-bot commented Feb 7, 2025 •

edited

Loading

flyte-bot commented Feb 7, 2025 •

edited

Loading

flyte-bot commented Feb 7, 2025 •

edited

Loading