From 99931c2c79229be078f88c8cee5abcdf9de9b8f5 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 4 Jan 2022 03:02:10 +0800 Subject: [PATCH 01/19] [pr into #785] Turn structured dataset into dataclass Signed-off-by: Kevin Su --- dataset_example.py | 12 +++-- flytekit/core/type_engine.py | 25 ++++++++++ flytekit/models/types.py | 2 +- flytekit/types/structured/basic_dfs.py | 3 +- .../types/structured/structured_dataset.py | 46 ++++++++++--------- tests/flytekit/unit/core/test_type_engine.py | 30 ++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 33 +++++++++++++ 7 files changed, 122 insertions(+), 29 deletions(-) diff --git a/dataset_example.py b/dataset_example.py index 74515b81bd..263edbbe90 100644 --- a/dataset_example.py +++ b/dataset_example.py @@ -22,8 +22,8 @@ ) from flytekit.types.structured.utils import get_filesystem -PANDAS_PATH = "/tmp/pandas" -NUMPY_PATH = "/tmp/numpy" +PANDAS_PATH = "s3://flyte-batch/my-s3-bucket/test-data/pandas" +NUMPY_PATH = "s3://flyte-batch/my-s3-bucket/test-data/numpy" BQ_PATH = "bq://photo-313016:flyte.new_table3" # https://github.com/flyteorg/flyte/issues/523 @@ -43,7 +43,7 @@ def t1(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, my_cols]: @task def t1a(dataframe: pd.DataFrame) -> StructuredDataset[my_cols, PARQUET]: # S3 (parquet) -> Pandas -> S3 (parquet) default behaviour - return StructuredDataset(dataframe=dataframe, file_format=PARQUET) + return StructuredDataset(dataframe=dataframe, file_format=PARQUET, uri=PANDAS_PATH) @task @@ -55,7 +55,7 @@ def t2(dataframe: pd.DataFrame) -> Annotated[pd.DataFrame, arrow_schema]: @task def t3(dataset: StructuredDataset[my_cols]) -> StructuredDataset[my_cols]: # s3 (parquet) -> pandas -> s3 (parquet) - print("Pandas dataframe") + print("Pandas dataframe:") print(dataset.open(pd.DataFrame).all()) # In the example, we download dataset when we open it. # Here we won't upload anything, since we're returning just the input object. @@ -91,7 +91,7 @@ def t6(dataset: StructuredDataset[my_cols]) -> pd.DataFrame: def t7(df1: pd.DataFrame, df2: pd.DataFrame) -> (StructuredDataset[my_cols], StructuredDataset[my_cols]): # df1: pandas -> bq # df2: pandas -> s3 (parquet) - return StructuredDataset(dataframe=df1, uri=BQ_PATH), StructuredDataset(dataframe=df1, file_format=PARQUET) + return StructuredDataset(dataframe=df1, uri=BQ_PATH), StructuredDataset(dataframe=df2, file_format=PARQUET) @task @@ -141,6 +141,8 @@ def decode( FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, "/", "parquet")) FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, "/", "parquet")) +FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, "s3://", "parquet")) +FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, "s3://", "parquet")) @task diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 9ef9057eba..25d201cbcc 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -42,6 +42,7 @@ Primitive, Scalar, Schema, + StructuredDatasetMetadata, ) from flytekit.models.types import LiteralType, SimpleType @@ -280,6 +281,7 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import FlyteFile from flytekit.types.schema.types import FlyteSchema + from flytekit.types.structured.structured_dataset import StructuredDataset for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) @@ -288,6 +290,7 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): issubclass(field_type, FlyteSchema) or issubclass(field_type, FlyteFile) or issubclass(field_type, FlyteDirectory) + or issubclass(field_type, StructuredDataset) ): lv = TypeEngine.to_literal(FlyteContext.current_context(), v, field_type, None) # dataclass_json package will extract the "path" from FlyteFile, FlyteDirectory, and write it to a @@ -300,6 +303,14 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): # as determined by the transformer. if issubclass(field_type, FlyteFile) or issubclass(field_type, FlyteDirectory): python_val.__setattr__(f.name, field_type(path=lv.scalar.blob.uri)) + elif issubclass(field_type, StructuredDataset): + python_val.__setattr__( + f.name, + field_type( + uri=lv.scalar.structured_dataset.uri, + file_format=lv.scalar.structured_dataset.metadata.format, + ), + ) elif dataclasses.is_dataclass(field_type): self._serialize_flyte_type(v, field_type) @@ -308,6 +319,7 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + from flytekit.types.structured.structured_dataset import StructuredDataset, StructuredDatasetTransformerEngine if not dataclasses.is_dataclass(expected_python_type): return python_val @@ -353,6 +365,19 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> ), expected_python_type, ) + elif issubclass(expected_python_type, StructuredDataset): + return StructuredDatasetTransformerEngine().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + structured_dataset=StructuredDataset( + metadata=StructuredDatasetMetadata(format=python_val.file_format), + uri=python_val.uri, + ) + ) + ), + expected_python_type, + ) else: for f in dataclasses.fields(expected_python_type): value = python_val.__getattribute__(f.name) diff --git a/flytekit/models/types.py b/flytekit/models/types.py index 0a82386164..0cc6092d22 100644 --- a/flytekit/models/types.py +++ b/flytekit/models/types.py @@ -161,7 +161,7 @@ def external_schema_bytes(self) -> bytes: def to_flyte_idl(self) -> _types_pb2.StructuredDatasetType: return _types_pb2.StructuredDatasetType( columns=[c.to_flyte_idl() for c in self.columns] if self.columns else None, - format=self._format, + format=self.format, external_schema_type=self.external_schema_type if self.external_schema_type else None, external_schema_bytes=self.external_schema_bytes if self.external_schema_bytes else None, ) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index dadad92cbc..00a47209fc 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -17,6 +17,7 @@ FLYTE_DATASET_TRANSFORMER, LOCAL, PARQUET, + S3, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, @@ -116,7 +117,7 @@ def decode( return spark.read.parquet(path) -for protocol in [LOCAL]: # Think how to add S3 and GCS +for protocol in [LOCAL, S3]: # Think how to add S3 and GCS FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(pa.Table, protocol, PARQUET)) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index fe00a5b17a..dc2fc68109 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -1,12 +1,18 @@ from __future__ import annotations import collections +import inspect +import os import re import types import typing from abc import ABC, abstractmethod +from dataclasses import dataclass, field from typing import Dict, Generator, Optional, Type, Union +from dataclasses_json import config, dataclass_json +from marshmallow import fields + try: from typing import Annotated, get_args, get_origin except ImportError: @@ -28,15 +34,19 @@ T = typing.TypeVar("T") # StructuredDataset type or a dataframe type DF = typing.TypeVar("DF") # Dataframe type -BIGQUERY = "bq" -S3 = "s3" +BIGQUERY = "bq://" +S3 = "s3://" LOCAL = "/" # Dataset File format PARQUET = "parquet" +@dataclass_json +@dataclass class StructuredDataset(object): + uri: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) + file_format: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) """ This is the user facing StructuredDataset class. Please don't confuse it with the literals.StructuredDataset class (that is just a model, a Python class representation of the protobuf). @@ -95,8 +105,10 @@ def __init__( metadata: typing.Optional[literals.StructuredDatasetMetadata] = None, ): self._dataframe = dataframe - self._uri = uri - self._file_format = file_format + # Make these fields public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.uri = uri + self.file_format = file_format # This is a special attribute that indicates if the data was either downloaded or uploaded self._metadata = metadata # This is not for users to set, the transformer will set this. @@ -108,18 +120,6 @@ def __init__( def dataframe(self) -> Type[typing.Any]: return self._dataframe - @property - def uri(self) -> Optional[str]: - return self._uri - - @uri.setter - def uri(self, uri: str): - self._uri = uri - - @classmethod - def file_format(cls) -> str: - return "" - @property def metadata(self) -> Optional[StructuredDatasetMetadata]: return self._metadata @@ -248,7 +248,7 @@ def decode( def protocol_prefix(uri: str) -> str: g = re.search(r"([\w]+)://.*", uri) if g and g.groups(): - return g.groups()[0] + return g.groups()[0] + "://" return LOCAL @@ -337,7 +337,7 @@ def to_literal( ) -> Literal: # If the type signature has the StructuredDataset class, it will, or at least should, also be a # StructuredDataset instance. - if issubclass(python_type, StructuredDataset): + if inspect.isclass(python_type) and issubclass(python_type, StructuredDataset): assert isinstance(python_val, StructuredDataset) # There are three cases that we need to take care of here. @@ -361,7 +361,7 @@ def to_literal( # Ex. # def t2(uri: str) -> StructuredDataset[my_cols] # return StructuredDataset(uri=uri) - format = python_val._file_format + format = python_val.file_format if python_val.dataframe is None: if not python_val.uri: raise ValueError(f"If dataframe is not specified, then the uri should be specified. {python_val}") @@ -380,7 +380,7 @@ def to_literal( protocol = self.DEFAULT_PROTOCOLS[df_type] else: protocol = protocol_prefix(python_val.uri) - format = python_val._file_format + format = python_val.file_format return self.encode( ctx, python_val, @@ -390,6 +390,8 @@ def to_literal( ) # Otherwise assume it's a dataframe instance. Wrap it with some defaults + if get_origin(python_type) is Annotated: + python_type = get_args(python_type)[0] fmt = self.DEFAULT_FORMATS[python_type] protocol = self.DEFAULT_PROTOCOLS[python_type] meta = StructuredDatasetMetadata(format=fmt, structured_dataset_type=expected.structured_dataset_type) @@ -486,11 +488,11 @@ def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any] external_schema_bytes=typing.cast(pa.lib.Schema, hint_args[0]).to_string().encode(), ) # 2. Fill in columns by checking for StructuredDataset metadata. For example, StructuredDataset[my_cols, parquet] - elif issubclass(t, StructuredDataset): + elif inspect.isclass(t) and issubclass(t, StructuredDataset): for k, v in t.columns().items(): lt = self._get_dataset_column_literal_type(v) converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) - return StructuredDatasetType(columns=converted_cols, format=t.file_format()) + return StructuredDatasetType(columns=converted_cols, format=str(t.file_format)) # 3. pd.Dataframe else: return StructuredDatasetType(columns=converted_cols, format=PARQUET) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 0b81763ad8..a73277b686 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -41,6 +41,7 @@ from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer from flytekit.types.schema import FlyteSchema +from flytekit.types.structured.structured_dataset import StructuredDataset def test_type_engine(): @@ -586,6 +587,35 @@ class TestFileStruct(object): assert o.b.c["hello"].path == ot.b.c["hello"].path +def test_structured_dataset_in_dataclass(): + df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @dataclass_json + @dataclass + class InnerDatasetStruct(object): + a: StructuredDataset + + @dataclass_json + @dataclass + class DatasetStruct(object): + a: StructuredDataset + b: InnerDatasetStruct + + sd = StructuredDataset(dataframe=df, file_format="parquet") + o = DatasetStruct(a=sd, b=InnerDatasetStruct(a=sd)) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(DatasetStruct) + lv = tf.to_literal(ctx, o, DatasetStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=DatasetStruct) + + assert_frame_equal(df, ot.a.open(pd.DataFrame).all()) + assert_frame_equal(df, ot.b.a.open(pd.DataFrame).all()) + assert "parquet" == ot.a.file_format + assert "parquet" == ot.b.a.file_format + + # Enums should have string values class Color(Enum): RED = "red" diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 3f7b4c901c..8b21470900 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -9,9 +9,11 @@ from enum import Enum import pandas +import pandas as pd import pytest from dataclasses_json import dataclass_json from google.protobuf.struct_pb2 import Struct +from pandas._testing import assert_frame_equal import flytekit from flytekit import ContainerTask, Secret, SQLTask, dynamic, kwtypes, map_task @@ -35,6 +37,7 @@ from flytekit.types.directory import FlyteDirectory, TensorboardLogs from flytekit.types.file import FlyteFile, PNGImageFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode +from flytekit.types.structured.structured_dataset import StructuredDataset serialization_settings = context_manager.SerializationSettings( project="proj", @@ -431,6 +434,36 @@ def wf(path: str) -> os.PathLike: assert "/tmp/flyte/" in wf(path="s3://somewhere").path +def test_structured_dataset_in_dataclass(): + df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) + + @dataclass_json + @dataclass + class InnerDatasetStruct(object): + a: StructuredDataset + + @dataclass_json + @dataclass + class DatasetStruct(object): + a: StructuredDataset + b: InnerDatasetStruct + + @task + def t1(path: str) -> DatasetStruct: + sd = StructuredDataset(dataframe=df, file_format="parquet", uri=path) + return DatasetStruct(a=sd, b=InnerDatasetStruct(a=sd)) + + @workflow + def wf(path: str) -> DatasetStruct: + return t1(path=path) + + res = wf(path="/tmp/somewhere") + assert "parquet" == res.a.file_format + assert "parquet" == res.b.a.file_format + assert_frame_equal(df, res.a.open(pd.DataFrame).all()) + assert_frame_equal(df, res.b.a.open(pd.DataFrame).all()) + + def test_wf1_with_map(): @task def t1(a: int) -> int: From 55f2e3ea6bd03605af3d1b18f4ef414550c67429 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 5 Jan 2022 01:55:56 +0800 Subject: [PATCH 02/19] Inspect class Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index dc2fc68109..616db7980d 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -423,7 +423,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: return self.open_as(ctx, sd_literal, df_type=expected_python_type) # Either a StructuredDataset type or some dataframe type. - if issubclass(expected_python_type, StructuredDataset): + if inspect.isclass(expected_python_type) and issubclass(expected_python_type, StructuredDataset): # Just save the literal for now. If in the future we find that we need the StructuredDataset type hint # type also, we can add it. sd = expected_python_type( From 2ab49a8dd8a9a16f43e47c6c43827bb6e801a30b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 5 Jan 2022 02:38:28 +0800 Subject: [PATCH 03/19] Inspect class Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 616db7980d..8be20d7a8d 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -438,6 +438,8 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # If the requested type was not a StructuredDataset, then it means it was a plain dataframe type, which means # we should do the opening/downloading and whatever else it might entail right now. No iteration option here. + if get_origin(expected_python_type) is Annotated: + expected_python_type = get_args(expected_python_type)[0] return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type) def open_as(self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF]) -> DF: From 3254f374d6411dbb1ae177262316ebe08480ffb1 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 5 Jan 2022 03:40:37 +0800 Subject: [PATCH 04/19] Resolved conflict Signed-off-by: Kevin Su --- flytekit/types/structured/__init__.py | 2 -- flytekit/types/structured/basic_dfs.py | 31 -------------------------- 2 files changed, 33 deletions(-) diff --git a/flytekit/types/structured/__init__.py b/flytekit/types/structured/__init__.py index ef19409943..da80c89016 100644 --- a/flytekit/types/structured/__init__.py +++ b/flytekit/types/structured/__init__.py @@ -5,8 +5,6 @@ PandasToParquetEncodingHandler, ParquetToArrowDecodingHandler, ParquetToPandasDecodingHandler, - ParquetToSparkDecodingHandler, - SparkToParquetEncodingHandler, ) try: diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 00a47209fc..732b311e7d 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -88,39 +88,8 @@ def decode( return pq.read_table(path, filesystem=get_filesystem(path)) -class SparkToParquetEncodingHandler(StructuredDatasetEncoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def encode( - self, - ctx: FlyteContext, - structured_dataset: StructuredDataset, - ) -> literals.StructuredDataset: - path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_path() - df = typing.cast(DataFrame, structured_dataset.dataframe) - df.write.parquet(path) - return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(format=PARQUET)) - - -class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def decode( - self, - ctx: FlyteContext, - flyte_value: literals.StructuredDataset, - ) -> DataFrame: - spark = SparkSession.builder.getOrCreate() - path = flyte_value.uri or ctx.file_access.get_random_remote_path() - return spark.read.parquet(path) - - for protocol in [LOCAL, S3]: # Think how to add S3 and GCS FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(pa.Table, protocol, PARQUET)) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToArrowDecodingHandler(pa.Table, protocol, PARQUET)) - FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) - FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) From 9f36800aee5e09832248659fa1f7aad345bcde91 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 5 Jan 2022 03:42:38 +0800 Subject: [PATCH 05/19] Resolved conflict Signed-off-by: Kevin Su --- flytekit/types/structured/basic_dfs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 732b311e7d..fcf4944658 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -88,7 +88,7 @@ def decode( return pq.read_table(path, filesystem=get_filesystem(path)) -for protocol in [LOCAL, S3]: # Think how to add S3 and GCS +for protocol in [LOCAL]: # Think how to add S3 and GCS FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(pa.Table, protocol, PARQUET)) From baf9e7709a4615f9b6579f2600b3af18b7e4bbaa Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 5 Jan 2022 03:43:09 +0800 Subject: [PATCH 06/19] Resolved conflict Signed-off-by: Kevin Su --- flytekit/types/structured/basic_dfs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index fcf4944658..732b311e7d 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -88,7 +88,7 @@ def decode( return pq.read_table(path, filesystem=get_filesystem(path)) -for protocol in [LOCAL]: # Think how to add S3 and GCS +for protocol in [LOCAL, S3]: # Think how to add S3 and GCS FLYTE_DATASET_TRANSFORMER.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToParquetEncodingHandler(pa.Table, protocol, PARQUET)) From 7fd01ff182345f4e0c6d28320c4d99443d700db1 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 7 Jan 2022 17:52:04 +0800 Subject: [PATCH 07/19] update Signed-off-by: Kevin Su --- dataset_example.py | 28 ------------------- flytekit/core/interface.py | 2 +- .../types/structured/structured_dataset.py | 2 +- .../flytekitplugins/spark/schema.py | 2 +- 4 files changed, 3 insertions(+), 31 deletions(-) diff --git a/dataset_example.py b/dataset_example.py index 263edbbe90..8b4f3b1236 100644 --- a/dataset_example.py +++ b/dataset_example.py @@ -6,8 +6,6 @@ import pandas as pd import pyarrow as pa import pyarrow.parquet as pq -import pyspark.sql.dataframe -from pyspark.sql import SparkSession from flytekit import FlyteContext, kwtypes, task, workflow from flytekit.models import literals @@ -158,17 +156,6 @@ def t10(dataset: StructuredDataset[my_cols]) -> np.ndarray: return np_array -@task -def t11(dataframe: pyspark.sql.dataframe.DataFrame) -> StructuredDataset[my_cols]: - return StructuredDataset(dataframe, file_format=PARQUET) - - -@task -def t12(dataset: StructuredDataset[my_cols]) -> pyspark.sql.dataframe.DataFrame: - spark_df = dataset.open(pyspark.sql.dataframe.DataFrame).all() - return spark_df - - @task def generate_pandas() -> pd.DataFrame: return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -184,24 +171,11 @@ def generate_arrow() -> pa.Table: return pa.Table.from_pandas(pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})) -@task -def generate_spark_dataframe() -> pyspark.sql.dataframe.DataFrame: - data = [ - {"Category": "A", "ID": 1, "Value": 121.44, "Truth": True}, - {"Category": "B", "ID": 2, "Value": 300.01, "Truth": False}, - {"Category": "C", "ID": 3, "Value": 10.99, "Truth": None}, - {"Category": "E", "ID": 4, "Value": 33.87, "Truth": True}, - ] - spark = SparkSession.builder.getOrCreate() - return spark.createDataFrame(data) - - @workflow() def wf(): df = generate_pandas() np_array = generate_numpy() arrow_df = generate_arrow() - spark_df = generate_spark_dataframe() t1(dataframe=df) t1a(dataframe=df) t2(dataframe=df) @@ -215,8 +189,6 @@ def wf(): t8a(dataframe=arrow_df) t9(dataframe=np_array) t10(dataset=StructuredDataset(uri=NUMPY_PATH, file_format=PARQUET)) - dataset = t11(dataframe=spark_df) - t12(dataset=dataset) if __name__ == "__main__": diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 216e337282..74dc883a00 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -277,7 +277,7 @@ def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] """ type_hints = typing.get_type_hints(fn) signature = inspect.signature(fn) - return_annotation = type_hints.get("return", None) + return_annotation = signature.return_annotation outputs = extract_return_annotation(return_annotation) for k, v in outputs.items(): diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 8be20d7a8d..f9c3706a68 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -494,7 +494,7 @@ def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any] for k, v in t.columns().items(): lt = self._get_dataset_column_literal_type(v) converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) - return StructuredDatasetType(columns=converted_cols, format=str(t.file_format)) + return StructuredDatasetType(columns=converted_cols, format=str(t.file_format())) # 3. pd.Dataframe else: return StructuredDatasetType(columns=converted_cols, format=PARQUET) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/schema.py b/plugins/flytekit-spark/flytekitplugins/spark/schema.py index 4de01bf2b7..50c6cb9638 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/schema.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/schema.py @@ -136,6 +136,6 @@ def decode( return user_ctx.spark_session.read.parquet(flyte_value.uri) -for protocol in ["/", "s3"]: +for protocol in ["/", "s3://"]: FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) From 77016600c17cc0669ccdb4f1f2bdd4ca4cfdd205 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 7 Jan 2022 15:56:06 -0800 Subject: [PATCH 08/19] tests Signed-off-by: Yee Hing Tong --- .../unit/core/hint_handling/__init__.py | 0 tests/flytekit/unit/core/hint_handling/a.py | 19 +++++++++++++ tests/flytekit/unit/core/hint_handling/b.py | 23 ++++++++++++++++ .../core/hint_handling/test_hint_handling.py | 27 +++++++++++++++++++ 4 files changed, 69 insertions(+) create mode 100644 tests/flytekit/unit/core/hint_handling/__init__.py create mode 100644 tests/flytekit/unit/core/hint_handling/a.py create mode 100644 tests/flytekit/unit/core/hint_handling/b.py create mode 100644 tests/flytekit/unit/core/hint_handling/test_hint_handling.py diff --git a/tests/flytekit/unit/core/hint_handling/__init__.py b/tests/flytekit/unit/core/hint_handling/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/core/hint_handling/a.py b/tests/flytekit/unit/core/hint_handling/a.py new file mode 100644 index 0000000000..2178164a20 --- /dev/null +++ b/tests/flytekit/unit/core/hint_handling/a.py @@ -0,0 +1,19 @@ +try: + from typing import Annotated, get_args, get_origin +except ImportError: + from typing_extensions import Annotated, get_origin, get_args + + +class AA(object): + ... + + +my_aa = Annotated[AA, "some annotation"] + + +def t1() -> AA: + return AA() + + +def t2() -> my_aa: + return my_aa() diff --git a/tests/flytekit/unit/core/hint_handling/b.py b/tests/flytekit/unit/core/hint_handling/b.py new file mode 100644 index 0000000000..1ea6f0276a --- /dev/null +++ b/tests/flytekit/unit/core/hint_handling/b.py @@ -0,0 +1,23 @@ +from __future__ import annotations +""" +There's no difference betweeen this file and a.py except that we have postponed evaluation +of annotations turned on here. +""" +try: + from typing import Annotated, get_args, get_origin +except ImportError: + from typing_extensions import Annotated, get_origin, get_args + + +from .a import AA + + +my_aa = Annotated[AA, "some annotation"] + + +def t1() -> AA: + return AA() + + +def t2() -> my_aa: + return my_aa() diff --git a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py new file mode 100644 index 0000000000..f67609bcb5 --- /dev/null +++ b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py @@ -0,0 +1,27 @@ +try: + from typing import Annotated, get_args, get_origin, get_type_hints +except ImportError: + from typing_extensions import Annotated, get_origin, get_args, get_type_hints + +import inspect + + +from .a import t1 as a_t1, t2 as a_t2 +from .b import t1 as b_t1, t2 as b_t2 + + +def printer(fn): + print(f"FN: {fn.__name__}") + type_hints = get_type_hints(fn) + print(f"Type hints {type_hints} return type hint {type_hints.get('return', None)}") + signature = inspect.signature(fn) + print(f"Inspect signature: {signature}") + print("-------") + + +def test_hinting(): + printer(a_t1) + printer(a_t2) + printer(b_t1) + printer(b_t2) + From 8b71946320ee4c025479e49b976154ffcdbf57d6 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 7 Jan 2022 15:59:02 -0800 Subject: [PATCH 09/19] nit Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/hint_handling/test_hint_handling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py index f67609bcb5..d325b92295 100644 --- a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py +++ b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py @@ -11,7 +11,7 @@ def printer(fn): - print(f"FN: {fn.__name__}") + print(f"In {fn.__module__} FN: {fn.__name__}") type_hints = get_type_hints(fn) print(f"Type hints {type_hints} return type hint {type_hints.get('return', None)}") signature = inspect.signature(fn) From db6669435dc7749a84ea4bc55cd8b1b07af14c79 Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Fri, 7 Jan 2022 19:17:41 -0800 Subject: [PATCH 10/19] improve printing Signed-off-by: Yee Hing Tong --- tests/flytekit/unit/core/hint_handling/a.py | 4 ++-- tests/flytekit/unit/core/hint_handling/b.py | 4 ++-- .../unit/core/hint_handling/test_hint_handling.py | 12 ++++++++---- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/flytekit/unit/core/hint_handling/a.py b/tests/flytekit/unit/core/hint_handling/a.py index 2178164a20..e9726fcd61 100644 --- a/tests/flytekit/unit/core/hint_handling/a.py +++ b/tests/flytekit/unit/core/hint_handling/a.py @@ -11,9 +11,9 @@ class AA(object): my_aa = Annotated[AA, "some annotation"] -def t1() -> AA: +def t1(in1: int) -> AA: return AA() -def t2() -> my_aa: +def t2(in1: int) -> my_aa: return my_aa() diff --git a/tests/flytekit/unit/core/hint_handling/b.py b/tests/flytekit/unit/core/hint_handling/b.py index 1ea6f0276a..4f68dbed9b 100644 --- a/tests/flytekit/unit/core/hint_handling/b.py +++ b/tests/flytekit/unit/core/hint_handling/b.py @@ -15,9 +15,9 @@ my_aa = Annotated[AA, "some annotation"] -def t1() -> AA: +def t1(in1: int) -> AA: return AA() -def t2() -> my_aa: +def t2(in1: int) -> my_aa: return my_aa() diff --git a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py index d325b92295..1ca103591e 100644 --- a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py +++ b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py @@ -12,11 +12,15 @@ def printer(fn): print(f"In {fn.__module__} FN: {fn.__name__}") - type_hints = get_type_hints(fn) - print(f"Type hints {type_hints} return type hint {type_hints.get('return', None)}") + type_hints = get_type_hints(fn, include_extras=True) + print(f" Type hints {type_hints} return type hint") signature = inspect.signature(fn) - print(f"Inspect signature: {signature}") - print("-------") + print(f" Inspect signature: {signature}") + print("") + return_type = type_hints.get('return', None) + print(f" Origin: {get_origin(return_type)}") + print(f" Args: {get_args(return_type)}") + print("===============================\n") def test_hinting(): From 6b4f8b4ae18ed336c262a2a9cdfa4b0217aa954d Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 10 Jan 2022 12:50:55 +0800 Subject: [PATCH 11/19] Fixed tests and lint Signed-off-by: Kevin Su --- dataset_example.py | 13 ------------- flytekit/core/interface.py | 4 ++-- flytekit/core/type_engine.py | 7 ++++--- flytekit/types/structured/structured_dataset.py | 17 ++++++++--------- tests/flytekit/unit/core/hint_handling/a.py | 4 ++-- tests/flytekit/unit/core/hint_handling/b.py | 7 +++---- .../core/hint_handling/test_hint_handling.py | 15 ++++++++------- tests/flytekit/unit/core/test_interface.py | 15 +++++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 2 +- 9 files changed, 43 insertions(+), 41 deletions(-) diff --git a/dataset_example.py b/dataset_example.py index 31599346f2..4e034fc226 100644 --- a/dataset_example.py +++ b/dataset_example.py @@ -163,17 +163,6 @@ def t10(dataset: StructuredDataset[my_cols]) -> np.ndarray: return np_array -@task -def t11(dataframe: pyspark.sql.dataframe.DataFrame) -> StructuredDataset[my_cols]: - return StructuredDataset(dataframe) - - -@task -def t12(dataset: StructuredDataset[my_cols]) -> pyspark.sql.dataframe.DataFrame: - spark_df = dataset.open(pyspark.sql.dataframe.DataFrame).all() - return spark_df - - @task def generate_pandas() -> pd.DataFrame: return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) @@ -207,8 +196,6 @@ def wf(): t8a(dataframe=arrow_df) t9(dataframe=np_array) t10(dataset=StructuredDataset(uri=NUMPY_PATH)) - dataset = t11(dataframe=spark_df) - t12(dataset=dataset) if __name__ == "__main__": diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 74dc883a00..490fd6be86 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -275,9 +275,9 @@ def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] For now the fancy object, maybe in the future a dumb object. """ - type_hints = typing.get_type_hints(fn) + type_hints = typing.get_type_hints(fn, include_extras=True) signature = inspect.signature(fn) - return_annotation = signature.return_annotation + return_annotation = type_hints.get("return", None) outputs = extract_return_annotation(return_annotation) for k, v in outputs.items(): diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 25d201cbcc..9cca58b4bc 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -44,7 +44,7 @@ Schema, StructuredDatasetMetadata, ) -from flytekit.models.types import LiteralType, SimpleType +from flytekit.models.types import LiteralType, SimpleType, StructuredDatasetType T = typing.TypeVar("T") DEFINITIONS = "definitions" @@ -308,7 +308,6 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): f.name, field_type( uri=lv.scalar.structured_dataset.uri, - file_format=lv.scalar.structured_dataset.metadata.format, ), ) @@ -371,7 +370,9 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> Literal( scalar=Scalar( structured_dataset=StructuredDataset( - metadata=StructuredDatasetMetadata(format=python_val.file_format), + metadata=StructuredDatasetMetadata( + structured_dataset_type=StructuredDatasetType(format=python_val.file_format) + ), uri=python_val.uri, ) ) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index a83c8055ef..2bde52b758 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -47,7 +47,7 @@ @dataclass class StructuredDataset(object): uri: typing.Optional[os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) - file_format: typing.Optional[str] = field(default=None, metadata=config(mm_field=fields.String())) + file_format: typing.Optional[str] = field(default=PARQUET, metadata=config(mm_field=fields.String())) """ This is the user facing StructuredDataset class. Please don't confuse it with the literals.StructuredDataset class (that is just a model, a Python class representation of the protobuf). @@ -103,6 +103,7 @@ def __init__( dataframe: typing.Optional[typing.Any] = None, uri: Optional[str] = None, metadata: typing.Optional[literals.StructuredDatasetMetadata] = None, + **kwargs, ): self._dataframe = dataframe # Make these fields public, so that the dataclass transformer can set a value for it @@ -257,7 +258,7 @@ def decode( def protocol_prefix(uri: str) -> str: g = re.search(r"([\w]+)://.*", uri) if g and g.groups(): - return g.groups()[0] + "://" + return g.groups()[0] return LOCAL @@ -311,11 +312,9 @@ def _finder(self, handler_map, df_type: Type, protocol: str, format: str): raise ValueError(f"Failed to find a handler for {df_type}, protocol {protocol}, fmt {format}") def get_encoder(self, df_type: Type, protocol: str, format: str): - protocol.replace("://", "") return self._finder(self.ENCODERS, df_type, protocol, format) def get_decoder(self, df_type: Type, protocol: str, format: str): - protocol.replace("://", "") return self._finder(self.DECODERS, df_type, protocol, format) def _handler_finder(self, h: Handlers) -> Dict[str, Handlers]: @@ -370,8 +369,11 @@ def to_literal( ) -> Literal: # Make a copy in case we need to hand off to encoders, since we can't be sure of mutations. # Check first to see if it's even an SD type. For backwards compatibility, we may be getting a + if get_origin(python_type) is Annotated: + python_type = get_args(python_type)[0] sdt = StructuredDatasetType(format=self.DEFAULT_FORMATS.get(python_type, None)) - if expected.structured_dataset_type: + + if expected and expected.structured_dataset_type: sdt = StructuredDatasetType( columns=expected.structured_dataset_type.columns, format=expected.structured_dataset_type.format, @@ -379,9 +381,6 @@ def to_literal( external_schema_bytes=expected.structured_dataset_type.external_schema_bytes, ) - if get_origin(python_type) is Annotated: - python_type = get_args(python_type)[0] - # If the type signature has the StructuredDataset class, it will, or at least should, also be a # StructuredDataset instance. if inspect.isclass(python_type) and issubclass(python_type, StructuredDataset): @@ -438,7 +437,7 @@ def to_literal( python_type = get_args(python_type)[0] fmt = self.DEFAULT_FORMATS[python_type] protocol = self.DEFAULT_PROTOCOLS[python_type] - meta = StructuredDatasetMetadata(structured_dataset_type=expected.structured_dataset_type) + meta = StructuredDatasetMetadata(structured_dataset_type=expected.structured_dataset_type if expected else None) sd = StructuredDataset(dataframe=python_val, metadata=meta) return self.encode(ctx, sd, python_type, protocol, fmt, sdt) diff --git a/tests/flytekit/unit/core/hint_handling/a.py b/tests/flytekit/unit/core/hint_handling/a.py index e9726fcd61..bfa5674930 100644 --- a/tests/flytekit/unit/core/hint_handling/a.py +++ b/tests/flytekit/unit/core/hint_handling/a.py @@ -1,7 +1,7 @@ try: - from typing import Annotated, get_args, get_origin + from typing import Annotated except ImportError: - from typing_extensions import Annotated, get_origin, get_args + from typing_extensions import Annotated class AA(object): diff --git a/tests/flytekit/unit/core/hint_handling/b.py b/tests/flytekit/unit/core/hint_handling/b.py index 4f68dbed9b..3db472b069 100644 --- a/tests/flytekit/unit/core/hint_handling/b.py +++ b/tests/flytekit/unit/core/hint_handling/b.py @@ -1,17 +1,16 @@ from __future__ import annotations + """ There's no difference betweeen this file and a.py except that we have postponed evaluation of annotations turned on here. """ try: - from typing import Annotated, get_args, get_origin + from typing import Annotated except ImportError: - from typing_extensions import Annotated, get_origin, get_args - + from typing_extensions import Annotated from .a import AA - my_aa = Annotated[AA, "some annotation"] diff --git a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py index 1ca103591e..77edbf5c12 100644 --- a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py +++ b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py @@ -1,13 +1,14 @@ try: - from typing import Annotated, get_args, get_origin, get_type_hints + from typing import get_args, get_origin, get_type_hints except ImportError: - from typing_extensions import Annotated, get_origin, get_args, get_type_hints + from typing_extensions import get_origin, get_args, get_type_hints import inspect - -from .a import t1 as a_t1, t2 as a_t2 -from .b import t1 as b_t1, t2 as b_t2 +from .a import t1 as a_t1 +from .a import t2 as a_t2 +from .b import t1 as b_t1 +from .b import t2 as b_t2 def printer(fn): @@ -16,8 +17,9 @@ def printer(fn): print(f" Type hints {type_hints} return type hint") signature = inspect.signature(fn) print(f" Inspect signature: {signature}") + print(f" Inspect signature return_annotation: {signature.return_annotation}") print("") - return_type = type_hints.get('return', None) + return_type = type_hints.get("return", None) print(f" Origin: {get_origin(return_type)}") print(f" Args: {get_args(return_type)}") print("===============================\n") @@ -28,4 +30,3 @@ def test_hinting(): printer(a_t2) printer(b_t1) printer(b_t2) - diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index 8e55ee1bb4..3fd4a84325 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -15,6 +15,11 @@ from flytekit.types.file import FlyteFile from flytekit.types.pickle import FlytePickle +try: + from typing import Annotated +except ImportError: + from typing_extensions import Annotated + def test_extract_only(): def x() -> typing.NamedTuple("NT1", x_str=str, y_int=int): @@ -184,6 +189,16 @@ def z(a: int = 7, b: str = "eleven") -> typing.Tuple[int, str]: assert not params.parameters["b"].required assert params.parameters["b"].default.scalar.primitive.string_value == "eleven" + def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"]: + return a + + our_interface = transform_function_to_interface(z) + params = transform_inputs_to_parameters(ctx, our_interface) + assert params.parameters["a"].required + assert params.parameters["a"].default is None + assert our_interface.inputs == {"a": typing.Annotated[int, "some annotation"]} + assert our_interface.outputs == {"o0": typing.Annotated[int, "some annotation"]} + def test_parameters_with_docstring(): ctx = context_manager.FlyteContext.current_context() diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 8b21470900..54c1c77b8c 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -450,7 +450,7 @@ class DatasetStruct(object): @task def t1(path: str) -> DatasetStruct: - sd = StructuredDataset(dataframe=df, file_format="parquet", uri=path) + sd = StructuredDataset(dataframe=df, uri=path) return DatasetStruct(a=sd, b=InnerDatasetStruct(a=sd)) @workflow From 68daa48fdb780f278cabe1ba6e6f6b306061547e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 10 Jan 2022 13:16:31 +0800 Subject: [PATCH 12/19] Fixed tests Signed-off-by: Kevin Su --- flytekit/core/interface.py | 5 ++++- plugins/flytekit-spark/flytekitplugins/spark/schema.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 490fd6be86..8365e34b0d 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -275,7 +275,10 @@ def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] For now the fancy object, maybe in the future a dumb object. """ - type_hints = typing.get_type_hints(fn, include_extras=True) + try: + type_hints = typing.get_type_hints(fn, include_extras=True) + except TypeError: + type_hints = typing.get_type_hints(fn) signature = inspect.signature(fn) return_annotation = type_hints.get("return", None) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/schema.py b/plugins/flytekit-spark/flytekitplugins/spark/schema.py index dcd49d91ea..72dad4fc92 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/schema.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/schema.py @@ -137,6 +137,6 @@ def decode( return user_ctx.spark_session.read.parquet(flyte_value.uri) -for protocol in ["/", "s3://"]: +for protocol in ["/", "s3"]: FLYTE_DATASET_TRANSFORMER.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True) FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True) From 92286bce8808b52eb22a528f28b60dfcac76840b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 10 Jan 2022 13:37:29 +0800 Subject: [PATCH 13/19] Fixed tests Signed-off-by: Kevin Su --- flytekit/core/interface.py | 1 + tests/flytekit/unit/core/hint_handling/test_hint_handling.py | 5 ++++- tests/flytekit/unit/core/test_interface.py | 4 ++-- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/flytekit/core/interface.py b/flytekit/core/interface.py index 8365e34b0d..263ac05fb7 100644 --- a/flytekit/core/interface.py +++ b/flytekit/core/interface.py @@ -276,6 +276,7 @@ def transform_function_to_interface(fn: Callable, docstring: Optional[Docstring] """ try: + # include_extras can only be used in python >= 3.9 type_hints = typing.get_type_hints(fn, include_extras=True) except TypeError: type_hints = typing.get_type_hints(fn) diff --git a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py index 77edbf5c12..e3c6b942d2 100644 --- a/tests/flytekit/unit/core/hint_handling/test_hint_handling.py +++ b/tests/flytekit/unit/core/hint_handling/test_hint_handling.py @@ -13,7 +13,10 @@ def printer(fn): print(f"In {fn.__module__} FN: {fn.__name__}") - type_hints = get_type_hints(fn, include_extras=True) + try: + type_hints = get_type_hints(fn, include_extras=True) + except TypeError: + type_hints = get_type_hints(fn) print(f" Type hints {type_hints} return type hint") signature = inspect.signature(fn) print(f" Inspect signature: {signature}") diff --git a/tests/flytekit/unit/core/test_interface.py b/tests/flytekit/unit/core/test_interface.py index 3fd4a84325..2317bab1c2 100644 --- a/tests/flytekit/unit/core/test_interface.py +++ b/tests/flytekit/unit/core/test_interface.py @@ -196,8 +196,8 @@ def z(a: Annotated[int, "some annotation"]) -> Annotated[int, "some annotation"] params = transform_inputs_to_parameters(ctx, our_interface) assert params.parameters["a"].required assert params.parameters["a"].default is None - assert our_interface.inputs == {"a": typing.Annotated[int, "some annotation"]} - assert our_interface.outputs == {"o0": typing.Annotated[int, "some annotation"]} + assert our_interface.inputs == {"a": Annotated[int, "some annotation"]} + assert our_interface.outputs == {"o0": Annotated[int, "some annotation"]} def test_parameters_with_docstring(): From 8b256b8bc6b69a08f14cd5bacc3818b4092e4ccb Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 10 Jan 2022 23:40:26 +0800 Subject: [PATCH 14/19] Pandas read local dir instead of file Signed-off-by: Kevin Su --- flytekit/types/structured/basic_dfs.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index 6ec7970ae1..aa6640534c 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -61,11 +61,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) - frames = [pandas.read_parquet(os.path.join(local_dir, f)) for f in os.listdir(local_dir)] - if len(frames) == 1: - return frames[0] - elif len(frames) > 1: - return pandas.concat(frames, copy=True) + return pd.read_parquet(local_dir) class ArrowToParquetEncodingHandler(StructuredDatasetEncoder): From b80295d395d8b7249041ad9f9c828be1bbf2c5b3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 11 Jan 2022 02:07:14 +0800 Subject: [PATCH 15/19] Remove duplicate code Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 2bde52b758..dee1d9cbe2 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -433,8 +433,6 @@ def to_literal( ) # Otherwise assume it's a dataframe instance. Wrap it with some defaults - if get_origin(python_type) is Annotated: - python_type = get_args(python_type)[0] fmt = self.DEFAULT_FORMATS[python_type] protocol = self.DEFAULT_PROTOCOLS[python_type] meta = StructuredDatasetMetadata(structured_dataset_type=expected.structured_dataset_type if expected else None) From 705d2cee7976e73b6392e73b8f8a6d0a5805fdd2 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 11 Jan 2022 03:08:19 +0800 Subject: [PATCH 16/19] support FlyteSchema -> StructuredDataset Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index dee1d9cbe2..9c8e46b23d 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -468,15 +468,18 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # The literal that we get in might be an old FlyteSchema. # We'll continue to support this for the time being. if lv.scalar.schema is not None: + sd = StructuredDataset() + sd_literal = literals.StructuredDataset( + uri=lv.scalar.schema.uri, + metadata=literals.StructuredDatasetMetadata( + # Dataframe will always be serialized to parquet file by FlyteSchema transformer + structured_dataset_type=StructuredDatasetType(format=PARQUET) + ), + ) + sd._literal_sd = sd_literal if issubclass(expected_python_type, StructuredDataset): - raise ValueError("We do not support FlyteSchema -> StructuredDataset transformations") + return sd else: - sd_literal = literals.StructuredDataset( - uri=lv.scalar.schema.uri, - metadata=literals.StructuredDatasetMetadata( - structured_dataset_type=StructuredDatasetType(format=self.DEFAULT_FORMATS[expected_python_type]) - ), - ) return self.open_as(ctx, sd_literal, df_type=expected_python_type) # Either a StructuredDataset type or some dataframe type. From 6ef85d555cdbad784975b14761e7af58e7d1759b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 11 Jan 2022 10:02:12 +0800 Subject: [PATCH 17/19] Remove inspect.isclass Signed-off-by: Kevin Su --- flytekit/types/structured/structured_dataset.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index 9c8e46b23d..a6661d1c47 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -383,7 +383,7 @@ def to_literal( # If the type signature has the StructuredDataset class, it will, or at least should, also be a # StructuredDataset instance. - if inspect.isclass(python_type) and issubclass(python_type, StructuredDataset): + if issubclass(python_type, StructuredDataset): assert isinstance(python_val, StructuredDataset) # There are three cases that we need to take care of here. @@ -467,6 +467,8 @@ def encode( def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T: # The literal that we get in might be an old FlyteSchema. # We'll continue to support this for the time being. + if get_origin(expected_python_type) is Annotated: + expected_python_type = get_args(expected_python_type)[0] if lv.scalar.schema is not None: sd = StructuredDataset() sd_literal = literals.StructuredDataset( @@ -483,7 +485,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: return self.open_as(ctx, sd_literal, df_type=expected_python_type) # Either a StructuredDataset type or some dataframe type. - if inspect.isclass(expected_python_type) and issubclass(expected_python_type, StructuredDataset): + if issubclass(expected_python_type, StructuredDataset): # Just save the literal for now. If in the future we find that we need the StructuredDataset type hint # type also, we can add it. sd = expected_python_type( @@ -497,8 +499,6 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # If the requested type was not a StructuredDataset, then it means it was a plain dataframe type, which means # we should do the opening/downloading and whatever else it might entail right now. No iteration option here. - if get_origin(expected_python_type) is Annotated: - expected_python_type = get_args(expected_python_type)[0] return self.open_as(ctx, lv.scalar.structured_dataset, df_type=expected_python_type) def open_as(self, ctx: FlyteContext, sd: literals.StructuredDataset, df_type: Type[DF]) -> DF: @@ -551,7 +551,7 @@ def _get_dataset_type(self, t: typing.Union[Type[StructuredDataset], typing.Any] raise ValueError(f"Unrecognized Annotated type for StructuredDataset {t}") # 2. Fill in columns by checking for StructuredDataset metadata. For example, StructuredDataset[my_cols, parquet] - elif inspect.isclass(t) and issubclass(t, StructuredDataset): + elif issubclass(t, StructuredDataset): for k, v in t.columns().items(): lt = self._get_dataset_column_literal_type(v) converted_cols.append(StructuredDatasetType.DatasetColumn(name=k, literal_type=lt)) From dad5e0e96a8c256e6297662034fbbb489e45e218 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 11 Jan 2022 17:26:34 +0800 Subject: [PATCH 18/19] Fix tests Signed-off-by: Kevin Su --- dev-requirements.in | 2 +- dev-requirements.txt | 14 ++------------ 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/dev-requirements.in b/dev-requirements.in index 9743d0fc20..c8eb2c0601 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -1,6 +1,6 @@ -c requirements.txt -git+git://github.com/flyteorg/pytest-flyte@main#egg=pytest-flyte +git+https://github.com/flyteorg/pytest-flyte@main#egg=pytest-flyte coverage[toml] joblib mock diff --git a/dev-requirements.txt b/dev-requirements.txt index 28c974c3ba..08d48d1988 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -2,7 +2,7 @@ # This file is autogenerated by pip-compile with python 3.9 # To update, run: # -# make dev-requirements.txt +# pip-compile dev-requirements.in # -e file:.#egg=flytekit # via @@ -77,7 +77,6 @@ cryptography==36.0.1 # via # -c requirements.txt # paramiko - # secretstorage dataclasses-json==0.5.6 # via # -c requirements.txt @@ -138,11 +137,6 @@ importlib-metadata==4.10.0 # keyring iniconfig==1.1.1 # via pytest -jeepney==0.7.1 - # via - # -c requirements.txt - # keyring - # secretstorage jinja2==3.0.3 # via # -c requirements.txt @@ -262,7 +256,7 @@ pytest==6.2.5 # pytest-flyte pytest-docker==0.10.3 # via pytest-flyte -pytest-flyte @ git+git://github.com/flyteorg/pytest-flyte@main +pytest-flyte @ git+https://github.com/flyteorg/pytest-flyte@main # via -r dev-requirements.in python-dateutil==2.8.1 # via @@ -315,10 +309,6 @@ retry==0.9.2 # via # -c requirements.txt # flytekit -secretstorage==3.3.1 - # via - # -c requirements.txt - # keyring six==1.16.0 # via # -c requirements.txt From f70e32d9770f06ccdeedac440d9a12f5f61d26db Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 11 Jan 2022 17:47:28 +0800 Subject: [PATCH 19/19] Fix lint Signed-off-by: Kevin Su --- .github/workflows/pythonbuild.yml | 2 +- flytekit/types/structured/structured_dataset.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index d77197717f..7d5604be60 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -148,7 +148,7 @@ jobs: python-version: 3.8 - name: Install dependencies run: | - python -m pip install --upgrade pip==21.2.4 setuptools wheel + python -m pip install --upgrade pip setuptools wheel pip install -r doc-requirements.txt - name: Build the documentation run: make -C docs html diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index a6661d1c47..2b4d651b2e 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -1,7 +1,6 @@ from __future__ import annotations import collections -import inspect import os import re import types