From eb1c75e4fe6875640560a1d6dbdfb377c6c8e494 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 24 Feb 2022 23:53:10 +0800 Subject: [PATCH 1/9] Revisit StructuredDatasetDecoder interface Signed-off-by: Kevin Su --- flytekit/models/literals.py | 2 +- flytekit/types/structured/basic_dfs.py | 14 +++--- flytekit/types/structured/bigquery.py | 19 ++++---- .../types/structured/structured_dataset.py | 22 +++++----- .../flytekitplugins/spark/sd_transformers.py | 4 ++ plugins/flytekit-spark/tests/test_wf.py | 43 +++++++++---------- 6 files changed, 52 insertions(+), 52 deletions(-) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index bc398ab7a7..a7bdf43153 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -549,7 +549,7 @@ def from_flyte_idl(cls, pb2_object): class StructuredDatasetMetadata(_common.FlyteIdlEntity): - def __init__(self, structured_dataset_type: StructuredDatasetType = None): + def __init__(self, structured_dataset_type: StructuredDatasetType): self._structured_dataset_type = structured_dataset_type @property diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 49b2f13ed9..7e5a3d5e89 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -55,14 +55,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pd.DataFrame: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pd.read_parquet(local_dir, columns=columns) return pd.read_parquet(local_dir) @@ -94,14 +93,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pa.Table: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pq.read_table(local_dir, columns=columns) return pq.read_table(local_dir) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index 923ea06e9e..280210ead2 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -11,7 +11,6 @@ from flytekit.models.types import StructuredDatasetType from flytekit.types.structured.structured_dataset import ( BIGQUERY, - DF, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, @@ -29,7 +28,9 @@ def _write_to_bq(structured_dataset: StructuredDataset): client.load_table_from_dataframe(df, table_id) -def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: +def _read_from_bq( + flyte_value: literals.StructuredDataset, current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None +) -> pd.DataFrame: path = flyte_value.uri _, project_id, dataset_id, table_id = re.split("\\.|://|:", path) client = bigquery_storage.BigQueryReadClient() @@ -37,10 +38,8 @@ def _read_from_bq(flyte_value: literals.StructuredDataset) -> pd.DataFrame: parent = "projects/{}".format(project_id) read_options = None - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] read_options = types.ReadSession.TableReadOptions(selected_fields=columns) requested_session = types.ReadSession(table=table, data_format=types.DataFormat.ARROW, read_options=read_options) @@ -78,8 +77,9 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - ) -> typing.Union[DF, typing.Generator[DF, None, None]]: - return _read_from_bq(flyte_value) + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + ) -> pd.DataFrame: + return _read_from_bq(flyte_value, current_task_metadata) class ArrowToBQEncodingHandlers(StructuredDatasetEncoder): @@ -106,7 +106,8 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - ) -> typing.Union[DF, typing.Generator[DF, None, None]]: + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + ) -> pa.Table: return pa.Table.from_pandas(_read_from_bq(flyte_value)) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 819bc012cc..11e816187a 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -261,14 +261,17 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: Optional[StructuredDatasetMetadata] = None, ) -> Union[DF, Generator[DF, None, None]]: """ This is code that will be called by the dataset transformer engine to ultimately translate from a Flyte Literal value into a Python instance. - :param ctx: + :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param flyte_value: This will be a Flyte IDL StructuredDataset Literal - do not confuse this with the StructuredDataset class defined also in this module. + :param current_task_metadata: Metadata contains column name and type, and decoder will use it to + read specific column of parquet file or database table .It might be different from the metadata in the incoming literal. :return: This function can either return an instance of the dataframe that this decoder handles, or an iterator of those dataframes. """ @@ -596,8 +599,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... if column_dict is None or len(column_dict) == 0: # but if it does, then we just copy it over if incoming_columns is not None and incoming_columns != []: - for c in incoming_columns: - final_dataset_columns.append(c) + final_dataset_columns = incoming_columns.copy() # If the current running task's input does have columns defined else: final_dataset_columns = self._convert_ordered_dict_of_columns_to_list(column_dict) @@ -614,7 +616,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... # t1(input_a: StructuredDataset) # or # t1(input_a: Annotated[StructuredDataset, my_cols]) if issubclass(expected_python_type, StructuredDataset): - sd = expected_python_type( + sd = StructuredDataset( dataframe=None, # Note here that the type being passed in metadata=metad, @@ -634,19 +636,15 @@ def open_as( updated_metadata: Optional[StructuredDatasetMetadata] = None, ) -> DF: """ - - :param ctx: + :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param sd: :param df_type: - :param meta: New metadata type, since it might be different from the metadata in the literal. - :return: + :param updated_metadata: New metadata type, since it might be different from the metadata in the literal. + :return: dataframe. It could be pandas dataframe or arrow table, etc. """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) - # todo: revisit this, we probably should add a new field to the decoder interface - if updated_metadata: - sd._metadata = updated_metadata - result = decoder.decode(ctx, sd) + result = decoder.decode(ctx, sd, updated_metadata) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index 2466d3fc13..7c9a2af0f5 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -39,8 +39,12 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> DataFrame: user_ctx = FlyteContext.current_context().user_space_params + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] + return user_ctx.spark_session.read.parquet(flyte_value.uri).select(*columns) return user_ctx.spark_session.read.parquet(flyte_value.uri) diff --git a/plugins/flytekit-spark/tests/test_wf.py b/plugins/flytekit-spark/tests/test_wf.py index a0a624fec7..8c42a6162f 100644 --- a/plugins/flytekit-spark/tests/test_wf.py +++ b/plugins/flytekit-spark/tests/test_wf.py @@ -6,6 +6,11 @@ from flytekit import kwtypes, task, workflow from flytekit.types.schema import FlyteSchema +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + def test_wf1_with_spark(): @task(task_config=Spark()) @@ -53,27 +58,6 @@ def my_wf() -> my_schema: assert df2 is not None -def test_ddwf1_with_spark(): - @task(task_config=Spark()) - def my_spark(a: int) -> (int, str): - session = flytekit.current_context().spark_session - assert session.sparkContext.appName == "FlyteSpark: ex:local:local:local" - return a + 2, "world" - - @task - def t2(a: str, b: str) -> str: - return b + a - - @workflow - def my_wf(a: int, b: str) -> (int, str): - x, y = my_spark(a=a) - d = t2(a=y, b=b) - return x, d - - x = my_wf(a=5, b="hello ") - assert x == (7, "hello world") - - def test_fs_sd_compatibility(): my_schema = FlyteSchema[kwtypes(name=str, age=int)] @@ -108,7 +92,6 @@ def test_spark_dataframe_return(): def my_spark(a: int) -> my_schema: session = flytekit.current_context().spark_session df = session.createDataFrame([("Alice", a)], my_schema.column_names()) - print(type(df)) return df @workflow @@ -120,3 +103,19 @@ def my_wf(a: int) -> my_schema: df2 = reader.all() result_df = df2.reset_index(drop=True) == pd.DataFrame(data={"name": ["Alice"], "age": [5]}).reset_index(drop=True) assert result_df.all().all() + + +def test_read_spark_subset_columns(): + @task + def t1() -> pd.DataFrame: + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @task(task_config=Spark()) + def t2(df: Annotated[pyspark.sql.DataFrame, kwtypes(Name=str)]) -> int: + return len(df.columns) + + @workflow() + def wf() -> int: + return t2(df=t1()) + + assert wf() == 1 From f08085a500d6d43107de15844f738672301f2880 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 00:44:29 +0800 Subject: [PATCH 2/9] Fixed tests Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 8 ++++---- .../flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py | 7 +++---- .../flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py | 7 +++---- tests/flytekit/unit/core/test_structured_dataset.py | 4 +++- .../test_structured_dataset_workflow.py | 2 ++ 5 files changed, 15 insertions(+), 13 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 11e816187a..db66214466 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -616,7 +616,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... # t1(input_a: StructuredDataset) # or # t1(input_a: Annotated[StructuredDataset, my_cols]) if issubclass(expected_python_type, StructuredDataset): - sd = StructuredDataset( + sd = expected_python_type( dataframe=None, # Note here that the type being passed in metadata=metad, @@ -633,18 +633,18 @@ def open_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - updated_metadata: Optional[StructuredDatasetMetadata] = None, + metadata: Optional[StructuredDatasetMetadata] = None, ) -> DF: """ :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param sd: :param df_type: - :param updated_metadata: New metadata type, since it might be different from the metadata in the literal. + :param metadata: New metadata type, since it might be different from the metadata in the literal. :return: dataframe. It could be pandas dataframe or arrow table, etc. """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) - result = decoder.decode(ctx, sd, updated_metadata) + result = decoder.decode(ctx, sd, metadata) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py index c17e2fc8bd..fdcae510ca 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py @@ -47,6 +47,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pa.Table: uri = flyte_value.uri if not ctx.file_access.is_remote(uri): @@ -54,10 +55,8 @@ def decode( _, path = split_protocol(uri) columns = None - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: fs = FSSpecPersistence.get_filesystem(uri) return pq.read_table(path, filesystem=fs, columns=columns) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py index 52bcc4522a..195b3d27b0 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py @@ -58,14 +58,13 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pd.DataFrame: uri = flyte_value.uri columns = None kwargs = get_storage_options(uri) - if flyte_value.metadata.structured_dataset_type.columns: - columns = [] - for c in flyte_value.metadata.structured_dataset_type.columns: - columns.append(c.name) + if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: return pd.read_parquet(uri, columns=columns, storage_options=kwargs) except NoCredentialsError: diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index a7ef1ea953..d64a683c19 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -221,6 +221,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]: yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -241,8 +242,9 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pd.DataFrame: - pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) StructuredDatasetTransformerEngine.register( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py index d911f971e4..49016960fe 100644 --- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py +++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py @@ -54,6 +54,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> pd.DataFrame: return pd_df @@ -86,6 +87,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, + current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, ) -> typing.Union[DF, typing.Generator[DF, None, None]]: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() From d425dbe7cc856aba698212958aaef93a09da52ed Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 01:26:20 +0800 Subject: [PATCH 3/9] Fixed tests Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index db66214466..6abc567a58 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -633,18 +633,18 @@ def open_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - metadata: Optional[StructuredDatasetMetadata] = None, + updated_metadata: Optional[StructuredDatasetMetadata] = None, ) -> DF: """ :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param sd: :param df_type: - :param metadata: New metadata type, since it might be different from the metadata in the literal. + :param updated_metadata: New metadata type, since it might be different from the metadata in the literal. :return: dataframe. It could be pandas dataframe or arrow table, etc. """ protocol = protocol_prefix(sd.uri) decoder = self.get_decoder(df_type, protocol, sd.metadata.structured_dataset_type.format) - result = decoder.decode(ctx, sd, metadata) + result = decoder.decode(ctx, sd, updated_metadata) if isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} returned iterator {result} but whole value requested from {sd}") return result From 2a3d5c526c262cdddd68d0a0aee396bd7652cd52 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 01:46:03 +0800 Subject: [PATCH 4/9] Fixed tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_structured_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index d64a683c19..e9282fb905 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -290,7 +290,7 @@ def test_to_python_value_with_incoming_columns(): # check when columns are not specified, should pull both and add column information. sd = fdt.to_python_value(ctx, lit, StructuredDataset) - assert sd.metadata.structured_dataset_type.columns[0].name == "age" + assert len(sd.metadata.structured_dataset_type.columns) == 2 # should also work if subset type is just an annotated pd.DataFrame subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)] From c1accfadc199a5245d14c49deed26693c28a1e6d Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 03:52:03 +0800 Subject: [PATCH 5/9] Address comment Signed-off-by: Kevin Su --- flytekit/models/literals.py | 3 +-- flytekit/types/structured/basic_dfs.py | 4 ++-- flytekit/types/structured/bigquery.py | 6 +++--- flytekit/types/structured/structured_dataset.py | 16 +++++++--------- .../flytekitplugins/fsspec/arrow.py | 2 +- .../flytekitplugins/fsspec/pandas.py | 2 +- .../flytekitplugins/spark/sd_transformers.py | 2 +- .../unit/core/test_structured_dataset.py | 4 ++-- .../test_structured_dataset_workflow.py | 4 ++-- 9 files changed, 20 insertions(+), 23 deletions(-) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index a7bdf43153..88ad43de55 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -1,4 +1,3 @@ -import typing from datetime import datetime as _datetime import pytz as _pytz @@ -571,7 +570,7 @@ def from_flyte_idl(cls, pb2_object: _literals_pb2.StructuredDatasetMetadata) -> class StructuredDataset(_common.FlyteIdlEntity): - def __init__(self, uri: str, metadata: typing.Optional[StructuredDatasetMetadata] = None): + def __init__(self, uri: str, metadata: StructuredDatasetMetadata): """ A strongly typed schema that defines the interface of data retrieved from the underlying storage medium. """ diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 7e5a3d5e89..0d07fda770 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -55,7 +55,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() @@ -93,7 +93,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pa.Table: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index 280210ead2..d0212c246b 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -29,7 +29,7 @@ def _write_to_bq(structured_dataset: StructuredDataset): def _read_from_bq( - flyte_value: literals.StructuredDataset, current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None + flyte_value: literals.StructuredDataset, current_task_metadata: StructuredDatasetMetadata ) -> pd.DataFrame: path = flyte_value.uri _, project_id, dataset_id, table_id = re.split("\\.|://|:", path) @@ -77,7 +77,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: return _read_from_bq(flyte_value, current_task_metadata) @@ -106,7 +106,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pa.Table: return pa.Table.from_pandas(_read_from_bq(flyte_value)) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 6abc567a58..eb379fe5e6 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -105,9 +105,7 @@ def all(self) -> DF: if self._dataframe_type is None: raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.") ctx = FlyteContextManager.current_context() - return flyte_dataset_transformer.open_as( - ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata - ) + return flyte_dataset_transformer.open_as(ctx, self.literal, self._dataframe_type, self.metadata) def iter(self) -> Generator[DF, None, None]: if self._dataframe_type is None: @@ -261,7 +259,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> Union[DF, Generator[DF, None, None]]: """ This is code that will be called by the dataset transformer engine to ultimately translate from a Flyte Literal @@ -270,8 +268,8 @@ def decode( :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes :param flyte_value: This will be a Flyte IDL StructuredDataset Literal - do not confuse this with the StructuredDataset class defined also in this module. - :param current_task_metadata: Metadata contains column name and type, and decoder will use it to - read specific column of parquet file or database table .It might be different from the metadata in the incoming literal. + :param current_task_metadata: Metadata object containing the type (and columns if any) for the currently + executing task. This type may have more or less information than the type information bundled inside the incoming flyte_value. :return: This function can either return an instance of the dataframe that this decoder handles, or an iterator of those dataframes. """ @@ -588,7 +586,7 @@ def t2(in_a: Annotated[StructuredDataset, kwtypes(col_b=float)]): ... sd._literal_sd = sd_literal return sd else: - return self.open_as(ctx, sd_literal, df_type=expected_python_type) + return self.open_as(ctx, sd_literal, expected_python_type, metad) # Start handling for StructuredDataset scalars, first look at the columns incoming_columns = lv.scalar.structured_dataset.metadata.structured_dataset_type.columns @@ -633,7 +631,7 @@ def open_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - updated_metadata: Optional[StructuredDatasetMetadata] = None, + updated_metadata: StructuredDatasetMetadata, ) -> DF: """ :param ctx: A FlyteContext, useful in accessing the filesystem and other attributes @@ -654,7 +652,7 @@ def iter_as( ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF], - updated_metadata: Optional[StructuredDatasetMetadata] = None, + updated_metadata: StructuredDatasetMetadata, ) -> Generator[DF, None, None]: protocol = protocol_prefix(sd.uri) decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py index fdcae510ca..45a23ed222 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py @@ -47,7 +47,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pa.Table: uri = flyte_value.uri if not ctx.file_access.is_remote(uri): diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py index 195b3d27b0..490cd5ef70 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py @@ -58,7 +58,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: uri = flyte_value.uri columns = None diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index 7c9a2af0f5..6a32457fe1 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -39,7 +39,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> DataFrame: user_ctx = FlyteContext.current_context().user_space_params if current_task_metadata and current_task_metadata.structured_dataset_type.columns: diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py index e9282fb905..8efaecfcc9 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -221,7 +221,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]: yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -242,7 +242,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py index 49016960fe..0fc0bd976c 100644 --- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py +++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py @@ -54,7 +54,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> pd.DataFrame: return pd_df @@ -87,7 +87,7 @@ def decode( self, ctx: FlyteContext, flyte_value: literals.StructuredDataset, - current_task_metadata: typing.Optional[StructuredDatasetMetadata] = None, + current_task_metadata: StructuredDatasetMetadata, ) -> typing.Union[DF, typing.Generator[DF, None, None]]: path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() From 44f2ce1f89c94d7d46357ba758846dc1d4c62928 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 03:54:00 +0800 Subject: [PATCH 6/9] Address comment Signed-off-by: Kevin Su --- flytekit/models/literals.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 88ad43de55..31c4bf0514 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -1,4 +1,5 @@ from datetime import datetime as _datetime +from typing import Optional import pytz as _pytz from flyteidl.core import literals_pb2 as _literals_pb2 @@ -548,7 +549,7 @@ def from_flyte_idl(cls, pb2_object): class StructuredDatasetMetadata(_common.FlyteIdlEntity): - def __init__(self, structured_dataset_type: StructuredDatasetType): + def __init__(self, structured_dataset_type: Optional[StructuredDatasetType] = None): self._structured_dataset_type = structured_dataset_type @property From 8292038f79de5f4a2f3deee91f41ede9f8de0447 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 17:33:54 +0800 Subject: [PATCH 7/9] Fixed tests Signed-off-by: Kevin Su --- flytekit/models/literals.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/models/literals.py b/flytekit/models/literals.py index 31c4bf0514..1a6e56cd90 100644 --- a/flytekit/models/literals.py +++ b/flytekit/models/literals.py @@ -571,7 +571,7 @@ def from_flyte_idl(cls, pb2_object: _literals_pb2.StructuredDatasetMetadata) -> class StructuredDataset(_common.FlyteIdlEntity): - def __init__(self, uri: str, metadata: StructuredDatasetMetadata): + def __init__(self, uri: str, metadata: Optional[StructuredDatasetMetadata] = None): """ A strongly typed schema that defines the interface of data retrieved from the underlying storage medium. """ From fe445a180374147ca2a665602fe096a8b1581277 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 25 Feb 2022 17:54:23 +0800 Subject: [PATCH 8/9] Fixed tests Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 5 +---- .../unit/core/test_structured_dataset_handlers.py | 10 +++++----- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index eb379fe5e6..fa57d7b5de 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -656,10 +656,7 @@ def iter_as( ) -> Generator[DF, None, None]: protocol = protocol_prefix(sd.uri) decoder = self.DECODERS[df_type][protocol][sd.metadata.structured_dataset_type.format] - # todo: revisit this, should we add a new field to the decoder interface - if updated_metadata: - sd._metadata = updated_metadata - result = decoder.decode(ctx, sd) + result = decoder.decode(ctx, sd, updated_metadata) if not isinstance(result, types.GeneratorType): raise ValueError(f"Decoder {decoder} didn't return iterator {result} but should have from {sd}") return result diff --git a/tests/flytekit/unit/core/test_structured_dataset_handlers.py b/tests/flytekit/unit/core/test_structured_dataset_handlers.py index 0d755ab78b..ada7483a0f 100644 --- a/tests/flytekit/unit/core/test_structured_dataset_handlers.py +++ b/tests/flytekit/unit/core/test_structured_dataset_handlers.py @@ -6,6 +6,7 @@ from flytekit.core import context_manager from flytekit.core.base_task import kwtypes +from flytekit.models.literals import StructuredDatasetMetadata from flytekit.models.types import StructuredDatasetType from flytekit.types.structured import basic_dfs from flytekit.types.structured.structured_dataset import ( @@ -26,12 +27,11 @@ def test_pandas(): decoder = basic_dfs.ParquetToPandasDecodingHandler("/") ctx = context_manager.FlyteContextManager.current_context() - sd = StructuredDataset( - dataframe=df, - ) - sd_lit = encoder.encode(ctx, sd, StructuredDatasetType(format="parquet")) + sd = StructuredDataset(dataframe=df) + sd_type = StructuredDatasetType(format="parquet") + sd_lit = encoder.encode(ctx, sd, sd_type) - df2 = decoder.decode(ctx, sd_lit) + df2 = decoder.decode(ctx, sd_lit, StructuredDatasetMetadata(sd_type)) assert df.equals(df2) From 245bf786c51005043e86271af2905ead5bc433d1 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 1 Mar 2022 18:24:31 +0800 Subject: [PATCH 9/9] Address comment Signed-off-by: Kevin Su --- flytekit/types/structured/basic_dfs.py | 4 ++-- flytekit/types/structured/bigquery.py | 2 +- plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py | 2 +- plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py | 2 +- .../flytekit-spark/flytekitplugins/spark/sd_transformers.py | 2 +- 5 files changed, 6 insertions(+), 6 deletions(-) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 0d07fda770..ff9d692cec 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -60,7 +60,7 @@ def decode( path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pd.read_parquet(local_dir, columns=columns) return pd.read_parquet(local_dir) @@ -98,7 +98,7 @@ def decode( path = flyte_value.uri local_dir = ctx.file_access.get_random_local_directory() ctx.file_access.get_data(path, local_dir, is_multipart=True) - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return pq.read_table(local_dir, columns=columns) return pq.read_table(local_dir) diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py index d0212c246b..aa0ef42f6b 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -38,7 +38,7 @@ def _read_from_bq( parent = "projects/{}".format(project_id) read_options = None - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] read_options = types.ReadSession.TableReadOptions(selected_fields=columns) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py index 45a23ed222..d47318666f 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/arrow.py @@ -55,7 +55,7 @@ def decode( _, path = split_protocol(uri) columns = None - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: fs = FSSpecPersistence.get_filesystem(uri) diff --git a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py index 490cd5ef70..07b58d243a 100644 --- a/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py +++ b/plugins/flytekit-data-fsspec/flytekitplugins/fsspec/pandas.py @@ -63,7 +63,7 @@ def decode( uri = flyte_value.uri columns = None kwargs = get_storage_options(uri) - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] try: return pd.read_parquet(uri, columns=columns, storage_options=kwargs) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index 6a32457fe1..cd451fa080 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -42,7 +42,7 @@ def decode( current_task_metadata: StructuredDatasetMetadata, ) -> DataFrame: user_ctx = FlyteContext.current_context().user_space_params - if current_task_metadata and current_task_metadata.structured_dataset_type.columns: + if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns: columns = [c.name for c in current_task_metadata.structured_dataset_type.columns] return user_ctx.spark_session.read.parquet(flyte_value.uri).select(*columns) return user_ctx.spark_session.read.parquet(flyte_value.uri)