[BUG] Support creation and reading of StructuredDataset with local or remote uri #2914

JiangJiaWei1103 · 2024-11-08T16:19:33Z

Tracking issue

Closes #5954.

Why are the changes needed?

When StructuredDataset is instantiated with an uri, the operation of directly reading the dataframe fails. There exist two cases to discuss:

Read the dataframe from StructuredDataset with a local file path as uri:

@task
def read_sd_from_local_uri() -> pd.DataFrame:
    sd = StructuredDataset(uri="./df.parquet", file_format="parquet")
    df = sd.open(pd.DataFrame).all()

    return df

Read the dataframe from StructuredDataset with a remote file path (e.g., s3 object storage) as uri:

@task
def read_sd_from_remote_uri() -> pd.DataFrame:
    sd = StructuredDataset(uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet", file_format="parquet")
    df = sd.open(pd.DataFrame).all()

    return df

Both cases should successfully read and return pd.DataFrame.

What changes were proposed in this pull request?

As commented here, users have no options to set _literal_sd in StructuredDataset. To prevent _literal_sd from being NoneType, we use StructuredDatasetTransformerEngine to construct a StructuredDataset literal and assign it to _literal_sd of StructuredDataset type.

How was this patch tested?

This patch is tested through a newly added unit test.

Setup process

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

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

Signed-off-by: JiaWei Jiang <[email protected]>

kumare3 · 2024-11-11T17:59:19Z

this is awesome

Future-Outlier

This is amazing, @JiangJiaWei1103
I am testing this.

Future-Outlier · 2024-11-12T14:11:07Z

flytekit/types/structured/structured_dataset.py

+    def _set_literal(self, ctx: FlyteContext, expected: LiteralType) -> None:
+        """
+        Explicitly set the StructuredDataset Literal to handle the following cases:
+
+        1. Read a dataframe from a StructuredDataset with an uri, for example:
+
+        @task
+        def return_sd() -> StructuredDataset:
+            sd = StructuredDataset(uri="s3://my-s3-bucket/s3_flyte_dir/df.parquet", file_format="parquet")
+            df = sd.open(pd.DataFrame).all()
+            return df
+
+        For details, please refer to this issue: https://github.com/flyteorg/flyte/issues/5954.
+        """


Very nice comment.

It means a lot to hear it from you!

Future-Outlier · 2024-11-12T14:12:58Z

flytekit/types/structured/structured_dataset.py

+        to_literal = loop_manager.synced(flyte_dataset_transformer.async_to_literal)
+        self._literal_sd = to_literal(ctx, self, StructuredDataset, expected).scalar.structured_dataset


Not sure if is here the best way to write it.
cc @wild-endeavor @thomasjpfan

Thanks for pointing this out. I'm also pondering if this is a good practice...

Will be glad to learn more from you guys!

Future-Outlier · 2024-11-12T14:18:02Z

tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py

+    @task
+    def upload_pqt_to_s3(local_path: str, remote_path: str) -> None:
+        """Upload local temp parquet file to s3 object storage"""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            fs = FileAccessProvider(
+                local_sandbox_dir=tmp_dir,
+                raw_output_prefix="s3://my-s3-bucket"
+            )
+            fs.upload(local_path, remote_path)
+
+    @task
+    def read_sd_from_uri(uri: str) -> pd.DataFrame:
+        sd = StructuredDataset(uri=uri, file_format="parquet")
+        df = sd.open(pd.DataFrame).all()
+
+        return df
+


Hi, @JiangJiaWei1103
you should put these tasks in a workflow,
which is closer to the reality that users use it.

No problem, let me modify the test snippets.

Signed-off-by: JiaWei Jiang <[email protected]>

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

Future-Outlier · 2024-11-14T14:18:48Z

tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py

+        ff = FlyteFile(path=REMOTE_PATH)
+        with ff.open(mode="rb") as f:
+            df_s3 = pd.read_parquet(f)


we shouldn't do this in a workflow, workflow is like a bridge to build your DAG.
task is the place to write Python code.

https://docs.flyte.org/en/latest/user_guide/concepts/main_concepts/workflows.html#divedeep-workflows

Future-Outlier · 2024-11-14T14:24:07Z

tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py

+    REMOTE_PATH = "s3://my-s3-bucket/my-test/df.parquet"
+    df = generate_pandas()
+


why REMOTE_PATH = "s3://my-s3-bucket/my-test/df.parquet" but not
REMOTE_PATH = "s3://my-s3-bucket/df.parquet"?

I forgot that I just created my-testdir and uploaded df.parquet manually last time. Hence, the test seems to pass without an error locally.

Your are right. We can just use REMOTE_PATH = "s3://my-s3-bucket/df.parquet". Also, the uploading logic in the previous comment should be modified, too.

yes the best way to test this in your computer.

(dev) future@outlier ~ % pwd /Users/future-outlier/code/dev/flytekit (dev) future@outlier ~ % pytest -s test_file.py

This is how flytekit run the uni test.
or you can just run

make unit_test

under the folder.

This is the expected result.

Signed-off-by: JiaWei Jiang <[email protected]>

codecov · 2024-11-14T15:00:54Z

Codecov Report

Attention: Patch coverage is 22.22222% with 7 lines in your changes missing coverage. Please review.

Project coverage is 75.81%. Comparing base (51f9a3e) to head (14df985).
Report is 8 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/types/structured/structured_dataset.py	22.22%	7 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2914       +/-   ##
===========================================
+ Coverage   46.91%   75.81%   +28.89%     
===========================================
  Files         199      199               
  Lines       20840    20860       +20     
  Branches     2681     2685        +4     
===========================================
+ Hits         9778    15815     +6037     
+ Misses      10582     4281     -6301     
- Partials      480      764      +284

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

Manually fill sd literal and metadata

8d3ea89

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 mentioned this pull request Nov 8, 2024

[BUG] Structured Dataset create with remote uri and read directly will fail flyteorg/flyte#5954

Open

2 tasks

JiangJiaWei1103 added 2 commits November 10, 2024 11:45

Use StructuredDatasetTransformerEngine to set sd literal

204d549

Signed-off-by: JiaWei Jiang <[email protected]>

Add unit test for reading sd from uri

56d9fd3

Signed-off-by: JiaWei Jiang <[email protected]>

JiangJiaWei1103 marked this pull request as ready for review November 10, 2024 06:43

JiangJiaWei1103 requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners November 10, 2024 06:43

JiangJiaWei1103 changed the title ~~[WIP] Support creation and reading of StructuredDataset with remote uri~~ Support creation and reading of StructuredDataset with local or remote uri Nov 10, 2024

JiangJiaWei1103 changed the title ~~Support creation and reading of StructuredDataset with local or remote uri~~ [BUG] Support creation and reading of StructuredDataset with local or remote uri Nov 10, 2024

Future-Outlier self-assigned this Nov 11, 2024

Future-Outlier reviewed Nov 12, 2024

View reviewed changes

JiangJiaWei1103 and others added 3 commits November 13, 2024 21:55

Put tasks into wf to mimic real-world use cases

e2344e9

Signed-off-by: JiaWei Jiang <[email protected]>

add env

2a689b8

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

env again

a30df5e

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

Future-Outlier reviewed Nov 14, 2024

View reviewed changes

JiangJiaWei1103 added 2 commits November 14, 2024 22:55

Modify python ff reading logic to a flyte task

72382ca

Signed-off-by: JiaWei Jiang <[email protected]>

Use task param instead of global path const

14df985

Signed-off-by: JiaWei Jiang <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BUG] Support creation and reading of StructuredDataset with local or remote uri #2914

[BUG] Support creation and reading of StructuredDataset with local or remote uri #2914

JiangJiaWei1103 commented Nov 8, 2024 •

edited

Loading

kumare3 commented Nov 11, 2024

Future-Outlier left a comment

Future-Outlier Nov 12, 2024

JiangJiaWei1103 Nov 13, 2024

Future-Outlier Nov 12, 2024

JiangJiaWei1103 Nov 13, 2024

Future-Outlier Nov 12, 2024

JiangJiaWei1103 Nov 13, 2024

Future-Outlier Nov 14, 2024

Future-Outlier Nov 14, 2024

JiangJiaWei1103 Nov 14, 2024 •

edited

Loading

Future-Outlier Nov 14, 2024

Future-Outlier Nov 14, 2024

codecov bot commented Nov 14, 2024 •

edited

Loading

		to_literal = loop_manager.synced(flyte_dataset_transformer.async_to_literal)
		self._literal_sd = to_literal(ctx, self, StructuredDataset, expected).scalar.structured_dataset

		REMOTE_PATH = "s3://my-s3-bucket/my-test/df.parquet"
		df = generate_pandas()

[BUG] Support creation and reading of StructuredDataset with local or remote uri #2914

Are you sure you want to change the base?

[BUG] Support creation and reading of StructuredDataset with local or remote uri #2914

Conversation

JiangJiaWei1103 commented Nov 8, 2024 • edited Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Check all the applicable boxes

Related PRs

Docs link

kumare3 commented Nov 11, 2024

Future-Outlier left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JiangJiaWei1103 Nov 14, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 14, 2024 • edited Loading

Codecov Report

JiangJiaWei1103 commented Nov 8, 2024 •

edited

Loading

JiangJiaWei1103 Nov 14, 2024 •

edited

Loading

codecov bot commented Nov 14, 2024 •

edited

Loading