From 99931c2c79229be078f88c8cee5abcdf9de9b8f5 Mon Sep 17 00:00:00 2001
From: Kevin Su <pingsutw@apache.org>
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 <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Wed, 5 Jan 2022 01:55:56 +0800
Subject: [PATCH 02/19] Inspect class

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Wed, 5 Jan 2022 02:38:28 +0800
Subject: [PATCH 03/19] Inspect class

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Wed, 5 Jan 2022 03:40:37 +0800
Subject: [PATCH 04/19] Resolved conflict

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Wed, 5 Jan 2022 03:42:38 +0800
Subject: [PATCH 05/19] Resolved conflict

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Wed, 5 Jan 2022 03:43:09 +0800
Subject: [PATCH 06/19] Resolved conflict

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Fri, 7 Jan 2022 17:52:04 +0800
Subject: [PATCH 07/19] update

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 7 Jan 2022 15:56:06 -0800
Subject: [PATCH 08/19] tests

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 .../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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 7 Jan 2022 15:59:02 -0800
Subject: [PATCH 09/19] nit

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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 <wild-endeavor@users.noreply.github.com>
Date: Fri, 7 Jan 2022 19:17:41 -0800
Subject: [PATCH 10/19] improve printing

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 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 <pingsutw@apache.org>
Date: Mon, 10 Jan 2022 12:50:55 +0800
Subject: [PATCH 11/19] Fixed tests and lint

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Mon, 10 Jan 2022 13:16:31 +0800
Subject: [PATCH 12/19] Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Mon, 10 Jan 2022 13:37:29 +0800
Subject: [PATCH 13/19] Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
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 <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Tue, 11 Jan 2022 02:07:14 +0800
Subject: [PATCH 15/19] Remove duplicate code

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Tue, 11 Jan 2022 03:08:19 +0800
Subject: [PATCH 16/19] support FlyteSchema -> StructuredDataset

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Tue, 11 Jan 2022 10:02:12 +0800
Subject: [PATCH 17/19] Remove inspect.isclass

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Tue, 11 Jan 2022 17:26:34 +0800
Subject: [PATCH 18/19] Fix tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 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 <pingsutw@apache.org>
Date: Tue, 11 Jan 2022 17:47:28 +0800
Subject: [PATCH 19/19] Fix lint

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 .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
