From a03bc8f98c8ed740ff0f8ed6f6879a7d3a05f88c Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 10:40:38 -0800
Subject: [PATCH 01/11] import singleton

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/extend/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/flytekit/extend/__init__.py b/flytekit/extend/__init__.py
index d420310fa2..8dfd68ac9a 100644
--- a/flytekit/extend/__init__.py
+++ b/flytekit/extend/__init__.py
@@ -46,3 +46,7 @@
 from flytekit.core.task import TaskPlugins
 from flytekit.core.type_engine import DictTransformer, T, TypeEngine, TypeTransformer
 from flytekit.tools.translator import get_serializable
+
+from flytekit.configuration.sdk import USE_STRUCTURED_DATASET
+if USE_STRUCTURED_DATASET.get():
+    from flytekit.types.structured.structured_dataset import FLYTE_DATASET_TRANSFORMER

From 30fed4f3405565e2bc84b1bdde5e9964cf6138cc Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 14:15:52 -0800
Subject: [PATCH 02/11] make compat a different job

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 .github/workflows/pythonbuild.yml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml
index bedf66383e..f43cb32c33 100644
--- a/.github/workflows/pythonbuild.yml
+++ b/.github/workflows/pythonbuild.yml
@@ -39,9 +39,11 @@ jobs:
         python -m pip install --upgrade pip==21.2.4 setuptools wheel
         make setup${{ matrix.spark-version-suffix }}
         pip freeze
+    - name: Test FlyteSchema compatibility
+      run: |
+        FLYTE_SDK_USE_STRUCTURED_DATASET=FALSE python -m pytest tests/flytekit_compatibility
     - name: Test with coverage
       run: |
-        coverage run -m pytest tests/flytekit_compatibility
         FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests/flytekit/unit
     - name: Integration Tests with coverage
       # https://github.com/actions/runner/issues/241#issuecomment-577360161

From d6be0eefe3c66a96008529160ab73091717f19f2 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 14:39:21 -0800
Subject: [PATCH 03/11] nit

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/types/structured/structured_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index ffefd101dd..45472d828f 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -723,8 +723,8 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
 
 
 if USE_STRUCTURED_DATASET.get():
-    logger.debug("Structured dataset module load... using structured datasets!")
+    logger.warning("Structured dataset module load... using structured datasets!")
     FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine()
     TypeEngine.register(FLYTE_DATASET_TRANSFORMER)
 else:
-    logger.debug("Structured dataset module load... not using structured datasets")
+    logger.warning("Structured dataset module load... not using structured datasets")

From 9603d86ff90a7e0a35cbdef0e10873f8d2c18410 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 15:57:04 -0800
Subject: [PATCH 04/11] wrong import location

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/types/structured/structured_dataset.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index 45472d828f..abfaa1c8d0 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -24,8 +24,7 @@
 
 from flytekit.configuration.sdk import USE_STRUCTURED_DATASET
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
-from flytekit.core.type_engine import TypeTransformer
-from flytekit.extend import TypeEngine
+from flytekit.core.type_engine import TypeTransformer, TypeEngine
 from flytekit.loggers import logger
 from flytekit.models import literals
 from flytekit.models import types as type_models

From 6daa7b540449b6c8752ae6d0e49fe08cf33b2926 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 16:53:03 -0800
Subject: [PATCH 05/11] turn certain functions into class methods

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/types/structured/basic_dfs.py        |  9 ++---
 flytekit/types/structured/bigquery.py         |  9 ++---
 .../types/structured/structured_dataset.py    | 35 ++++++++++++-------
 .../flytekitplugins/spark/sd_transformers.py  |  5 +--
 .../unit/core/test_structured_dataset.py      | 21 +++++------
 .../test_structured_dataset_workflow.py       |  9 ++---
 6 files changed, 51 insertions(+), 37 deletions(-)

diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py
index a65fbbeee3..f557172cc4 100644
--- a/flytekit/types/structured/basic_dfs.py
+++ b/flytekit/types/structured/basic_dfs.py
@@ -19,6 +19,7 @@
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
 )
 
 T = TypeVar("T")
@@ -107,7 +108,7 @@ def decode(
 
 
 for protocol in [LOCAL, S3]:  # Should we add 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(protocol), default_for_type=True)
-    FLYTE_DATASET_TRANSFORMER.register_handler(ParquetToArrowDecodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(ArrowToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(ParquetToArrowDecodingHandler(protocol), default_for_type=True)
diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py
index d9221e4110..b084aabe0e 100644
--- a/flytekit/types/structured/bigquery.py
+++ b/flytekit/types/structured/bigquery.py
@@ -17,6 +17,7 @@
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
     StructuredDatasetMetadata,
+    StructuredDatasetTransformerEngine,
 )
 
 
@@ -110,7 +111,7 @@ def decode(
         return pa.Table.from_pandas(_read_from_bq(flyte_value))
 
 
-FLYTE_DATASET_TRANSFORMER.register_handler(PandasToBQEncodingHandlers(), default_for_type=False)
-FLYTE_DATASET_TRANSFORMER.register_handler(BQToPandasDecodingHandler(), default_for_type=False)
-FLYTE_DATASET_TRANSFORMER.register_handler(ArrowToBQEncodingHandlers(), default_for_type=False)
-FLYTE_DATASET_TRANSFORMER.register_handler(BQToArrowDecodingHandler(), default_for_type=False)
+StructuredDatasetTransformerEngine.register_handler(PandasToBQEncodingHandlers(), default_for_type=False)
+StructuredDatasetTransformerEngine.register_handler(BQToPandasDecodingHandler(), default_for_type=False)
+StructuredDatasetTransformerEngine.register_handler(ArrowToBQEncodingHandlers(), default_for_type=False)
+StructuredDatasetTransformerEngine.register_handler(BQToArrowDecodingHandler(), default_for_type=False)
diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index abfaa1c8d0..86be59d823 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -333,10 +333,12 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
     DECODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetDecoder]]] = {}
     DEFAULT_PROTOCOLS: Dict[Type, str] = {}
     DEFAULT_FORMATS: Dict[Type, str] = {}
+    SINGLETON: Optional[StructuredDatasetTransformerEngine] = None
 
     Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder]
 
-    def _finder(self, handler_map, df_type: Type, protocol: str, format: str):
+    @staticmethod
+    def _finder(handler_map, df_type: Type, protocol: str, format: str):
         try:
             return handler_map[df_type][protocol][format]
         except KeyError:
@@ -351,18 +353,21 @@ 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):
-        return self._finder(self.ENCODERS, df_type, protocol, format)
+    @classmethod
+    def get_encoder(cls, df_type: Type, protocol: str, format: str):
+        return cls._finder(StructuredDatasetTransformerEngine.ENCODERS, df_type, protocol, format)
 
-    def get_decoder(self, df_type: Type, protocol: str, format: str):
-        return self._finder(self.DECODERS, df_type, protocol, format)
+    @classmethod
+    def get_decoder(cls, df_type: Type, protocol: str, format: str):
+        return cls._finder(StructuredDatasetTransformerEngine.DECODERS, df_type, protocol, format)
 
-    def _handler_finder(self, h: Handlers) -> Dict[str, Handlers]:
+    @classmethod
+    def _handler_finder(cls, h: Handlers) -> Dict[str, Handlers]:
         # Maybe think about default dict in the future, but is typing as nice?
         if isinstance(h, StructuredDatasetEncoder):
-            top_level = self.ENCODERS
+            top_level = cls.ENCODERS
         elif isinstance(h, StructuredDatasetDecoder):
-            top_level = self.DECODERS
+            top_level = cls.DECODERS
         else:
             raise TypeError(f"We don't support this type of handler {h}")
         if h.python_type not in top_level:
@@ -375,7 +380,8 @@ def __init__(self):
         super().__init__("StructuredDataset Transformer", StructuredDataset)
         self._type_assertions_enabled = False
 
-    def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False):
+    @classmethod
+    def register_handler(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False):
         """
         Call this with any handler to register it with this dataframe meta-transformer
 
@@ -385,7 +391,7 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True,
             logger.info(f"Structured datasets not enabled, not registering handler {h}")
             return
 
-        lowest_level = self._handler_finder(h)
+        lowest_level = cls._handler_finder(h)
         if h.supported_format in lowest_level and override is False:
             raise ValueError(f"Already registered a handler for {(h.python_type, h.protocol, h.supported_format)}")
         lowest_level[h.supported_format] = h
@@ -393,13 +399,15 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True,
 
         if default_for_type:
             # TODO: Add logging, think about better ux, maybe default False and warn if doesn't exist.
-            self.DEFAULT_FORMATS[h.python_type] = h.supported_format
-            self.DEFAULT_PROTOCOLS[h.python_type] = h.protocol
+            cls.DEFAULT_FORMATS[h.python_type] = h.supported_format
+            cls.DEFAULT_PROTOCOLS[h.python_type] = h.protocol
 
         # Register with the type engine as well
         # The semantics as of now are such that it doesn't matter which order these transformers are loaded in, as
         # long as the older Pandas/FlyteSchema transformer do not also specify the override
-        TypeEngine.register_additional_type(self, h.python_type, override=True)
+        if cls.SINGLETON is None:
+            raise ValueError(f"Attempting to register handler {h} before singleton was assigned")
+        TypeEngine.register_additional_type(cls.SINGLETON, h.python_type, override=True)
 
     def assert_type(self, t: Type[StructuredDataset], v: typing.Any):
         return
@@ -724,6 +732,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
 if USE_STRUCTURED_DATASET.get():
     logger.warning("Structured dataset module load... using structured datasets!")
     FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine()
+    StructuredDatasetTransformerEngine.SINGLETON = FLYTE_DATASET_TRANSFORMER
     TypeEngine.register(FLYTE_DATASET_TRANSFORMER)
 else:
     logger.warning("Structured dataset module load... not using structured datasets")
diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
index e0b1c7b41e..989b224ed8 100644
--- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
+++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
@@ -12,6 +12,7 @@
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
 )
 
 
@@ -45,5 +46,5 @@ def decode(
 
 
 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)
+    StructuredDatasetTransformerEngine.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True)
diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py
index 4e4309e292..107ec4a549 100644
--- a/tests/flytekit/unit/core/test_structured_dataset.py
+++ b/tests/flytekit/unit/core/test_structured_dataset.py
@@ -24,6 +24,7 @@
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
     convert_schema_type_to_structured_dataset_type,
     extract_cols_and_format,
     protocol_prefix,
@@ -117,10 +118,10 @@ def test_types_sd():
 
 
 def test_retrieving():
-    assert FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", PARQUET) is not None
+    assert StructuredDatasetTransformerEngine.get_encoder(pd.DataFrame, "/", PARQUET) is not None
     with pytest.raises(ValueError):
         # We don't have a default "" format encoder
-        FLYTE_DATASET_TRANSFORMER.get_encoder(pd.DataFrame, "/", "")
+        StructuredDatasetTransformerEngine.get_encoder(pd.DataFrame, "/", "")
 
     class TempEncoder(StructuredDatasetEncoder):
         def __init__(self, protocol):
@@ -129,15 +130,15 @@ def __init__(self, protocol):
         def encode(self):
             ...
 
-    FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs"), default_for_type=False)
+    StructuredDatasetTransformerEngine.register_handler(TempEncoder("gs"), default_for_type=False)
     with pytest.raises(ValueError):
-        FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs://"), default_for_type=False)
+        StructuredDatasetTransformerEngine.register_handler(TempEncoder("gs://"), default_for_type=False)
 
     class TempEncoder:
         pass
 
     with pytest.raises(TypeError, match="We don't support this type of handler"):
-        FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder, default_for_type=False)
+        StructuredDatasetTransformerEngine.register_handler(TempEncoder, default_for_type=False)
 
 
 def test_to_literal():
@@ -184,7 +185,7 @@ def encode(
         ) -> literals.StructuredDataset:
             return literals.StructuredDataset(uri="")
 
-    FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("myavro"), default_for_type=True)
+    StructuredDatasetTransformerEngine.register_handler(TempEncoder("myavro"), default_for_type=True)
     lt = TypeEngine.to_literal_type(MyDF)
     assert lt.structured_dataset_type.format == "myavro"
 
@@ -196,9 +197,9 @@ def encode(
 
     # Test that looking up encoders/decoders falls back to the "" encoder/decoder
     empty_format_temp_encoder = TempEncoder("")
-    FLYTE_DATASET_TRANSFORMER.register_handler(empty_format_temp_encoder, default_for_type=False)
+    StructuredDatasetTransformerEngine.register_handler(empty_format_temp_encoder, default_for_type=False)
 
-    res = FLYTE_DATASET_TRANSFORMER.get_encoder(MyDF, "tmpfs", "rando")
+    res = StructuredDatasetTransformerEngine.get_encoder(MyDF, "tmpfs", "rando")
     assert res is empty_format_temp_encoder
 
 
@@ -221,7 +222,7 @@ def decode(
         ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]:
             yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
 
-    FLYTE_DATASET_TRANSFORMER.register_handler(
+    StructuredDatasetTransformerEngine.register_handler(
         MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False
     )
     sd = StructuredDataset()
@@ -241,7 +242,7 @@ def decode(
         ) -> pd.DataFrame:
             pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
 
-    FLYTE_DATASET_TRANSFORMER.register_handler(
+    StructuredDatasetTransformerEngine.register_handler(
         MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True
     )
     sd = StructuredDataset()
diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
index fa879ff0d9..0a19f4fdf7 100644
--- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
+++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
@@ -19,6 +19,7 @@
     BIGQUERY,
     DF,
     FLYTE_DATASET_TRANSFORMER,
+    StructuredDatasetTransformerEngine,
     LOCAL,
     PARQUET,
     S3,
@@ -58,8 +59,8 @@ def decode(
         return pd_df
 
 
-FLYTE_DATASET_TRANSFORMER.register_handler(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True)
-FLYTE_DATASET_TRANSFORMER.register_handler(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True)
+StructuredDatasetTransformerEngine.register_handler(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True)
+StructuredDatasetTransformerEngine.register_handler(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True)
 
 
 class NumpyEncodingHandlers(StructuredDatasetEncoder):
@@ -95,8 +96,8 @@ def decode(
 
 
 for protocol in [LOCAL, S3]:
-    FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET))
-    FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET))
+    StructuredDatasetTransformerEngine.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET))
+    StructuredDatasetTransformerEngine.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET))
 
 
 @task

From 4657497f7becda6fba48fc43838be849b1759886 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 18:02:59 -0800
Subject: [PATCH 06/11] make class

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/__init__.py                          |  1 +
 flytekit/extend/__init__.py                   |  4 ---
 flytekit/types/structured/basic_dfs.py        |  1 -
 flytekit/types/structured/bigquery.py         |  1 -
 .../types/structured/structured_dataset.py    | 19 +++++-----
 .../flytekitplugins/spark/sd_transformers.py  |  1 -
 .../unit/core/test_structured_dataset.py      | 36 ++++++++++---------
 .../test_structured_dataset_workflow.py       |  1 -
 8 files changed, 29 insertions(+), 35 deletions(-)

diff --git a/flytekit/__init__.py b/flytekit/__init__.py
index 64c0a26530..303f34caa5 100644
--- a/flytekit/__init__.py
+++ b/flytekit/__init__.py
@@ -194,6 +194,7 @@
         StructuredDataset,
         StructuredDatasetFormat,
         StructuredDatasetType,
+        StructuredDatasetTransformerEngine,
     )
 
 __version__ = "0.0.0+develop"
diff --git a/flytekit/extend/__init__.py b/flytekit/extend/__init__.py
index 8dfd68ac9a..d420310fa2 100644
--- a/flytekit/extend/__init__.py
+++ b/flytekit/extend/__init__.py
@@ -46,7 +46,3 @@
 from flytekit.core.task import TaskPlugins
 from flytekit.core.type_engine import DictTransformer, T, TypeEngine, TypeTransformer
 from flytekit.tools.translator import get_serializable
-
-from flytekit.configuration.sdk import USE_STRUCTURED_DATASET
-if USE_STRUCTURED_DATASET.get():
-    from flytekit.types.structured.structured_dataset import FLYTE_DATASET_TRANSFORMER
diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py
index f557172cc4..21bc08796a 100644
--- a/flytekit/types/structured/basic_dfs.py
+++ b/flytekit/types/structured/basic_dfs.py
@@ -12,7 +12,6 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
-    FLYTE_DATASET_TRANSFORMER,
     LOCAL,
     PARQUET,
     S3,
diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py
index b084aabe0e..8f1eb53925 100644
--- a/flytekit/types/structured/bigquery.py
+++ b/flytekit/types/structured/bigquery.py
@@ -12,7 +12,6 @@
 from flytekit.types.structured.structured_dataset import (
     BIGQUERY,
     DF,
-    FLYTE_DATASET_TRANSFORMER,
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index 86be59d823..c2059207aa 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -105,7 +105,7 @@ def all(self) -> DF:
         if self._dataframe_type is None:
             raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.")
         ctx = FlyteContextManager.current_context()
-        return FLYTE_DATASET_TRANSFORMER.open_as(
+        return flyte_dataset_transformer.open_as(
             ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata
         )
 
@@ -113,7 +113,7 @@ def iter(self) -> Generator[DF, None, None]:
         if self._dataframe_type is None:
             raise ValueError("No dataframe type set. Use open() to set the local dataframe type you want to use.")
         ctx = FlyteContextManager.current_context()
-        return FLYTE_DATASET_TRANSFORMER.iter_as(
+        return flyte_dataset_transformer.iter_as(
             ctx, self.literal, self._dataframe_type, updated_metadata=self.metadata
         )
 
@@ -169,7 +169,7 @@ class StructuredDatasetEncoder(ABC):
     def __init__(self, python_type: Type[T], protocol: str, supported_format: Optional[str] = None):
         """
         Extend this abstract class, implement the encode function, and register your concrete class with the
-        FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle
+        StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle
         dataframe libraries. This is the encoding interface, meaning it is used when there is a Python value that the
         flytekit type engine is trying to convert into a Flyte Literal. For the other way, see
         the StructuredDatasetEncoder
@@ -229,7 +229,7 @@ class StructuredDatasetDecoder(ABC):
     def __init__(self, python_type: Type[DF], protocol: str, supported_format: Optional[str] = None):
         """
         Extend this abstract class, implement the decode function, and register your concrete class with the
-        FLYTE_DATASET_TRANSFORMER defined at this module level in order for the core flytekit type engine to handle
+        StructuredDatasetTransformerEngine class in order for the core flytekit type engine to handle
         dataframe libraries. This is the decoder interface, meaning it is used when there is a Flyte Literal value,
         and we have to get a Python value out of it. For the other way, see the StructuredDatasetEncoder
 
@@ -333,7 +333,6 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]):
     DECODERS: Dict[Type, Dict[str, Dict[str, StructuredDatasetDecoder]]] = {}
     DEFAULT_PROTOCOLS: Dict[Type, str] = {}
     DEFAULT_FORMATS: Dict[Type, str] = {}
-    SINGLETON: Optional[StructuredDatasetTransformerEngine] = None
 
     Handlers = Union[StructuredDatasetEncoder, StructuredDatasetDecoder]
 
@@ -405,9 +404,8 @@ def register_handler(cls, h: Handlers, default_for_type: Optional[bool] = True,
         # Register with the type engine as well
         # The semantics as of now are such that it doesn't matter which order these transformers are loaded in, as
         # long as the older Pandas/FlyteSchema transformer do not also specify the override
-        if cls.SINGLETON is None:
-            raise ValueError(f"Attempting to register handler {h} before singleton was assigned")
-        TypeEngine.register_additional_type(cls.SINGLETON, h.python_type, override=True)
+        engine = StructuredDatasetTransformerEngine()
+        TypeEngine.register_additional_type(engine, h.python_type, override=True)
 
     def assert_type(self, t: Type[StructuredDataset], v: typing.Any):
         return
@@ -731,8 +729,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
 
 if USE_STRUCTURED_DATASET.get():
     logger.warning("Structured dataset module load... using structured datasets!")
-    FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine()
-    StructuredDatasetTransformerEngine.SINGLETON = FLYTE_DATASET_TRANSFORMER
-    TypeEngine.register(FLYTE_DATASET_TRANSFORMER)
+    flyte_dataset_transformer = StructuredDatasetTransformerEngine()
+    TypeEngine.register(flyte_dataset_transformer)
 else:
     logger.warning("Structured dataset module load... not using structured datasets")
diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
index 989b224ed8..28e2ffdadc 100644
--- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
+++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
@@ -7,7 +7,6 @@
 from flytekit.models.literals import StructuredDatasetMetadata
 from flytekit.models.types import StructuredDatasetType
 from flytekit.types.structured.structured_dataset import (
-    FLYTE_DATASET_TRANSFORMER,
     PARQUET,
     StructuredDataset,
     StructuredDatasetDecoder,
diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py
index 107ec4a549..3d66e01b29 100644
--- a/tests/flytekit/unit/core/test_structured_dataset.py
+++ b/tests/flytekit/unit/core/test_structured_dataset.py
@@ -19,7 +19,6 @@
 
 from flytekit import kwtypes
 from flytekit.types.structured.structured_dataset import (
-    FLYTE_DATASET_TRANSFORMER,
     PARQUET,
     StructuredDataset,
     StructuredDatasetDecoder,
@@ -146,7 +145,9 @@ def test_to_literal():
     lt = TypeEngine.to_literal_type(pd.DataFrame)
     df = generate_pandas()
 
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
+    fdt = StructuredDatasetTransformerEngine()
+
+    lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
     assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET
     assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == PARQUET
 
@@ -154,16 +155,16 @@ def test_to_literal():
     sd_with_literal_and_df._literal_sd = lit
 
     with pytest.raises(ValueError, match="Shouldn't have specified both literal"):
-        FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_literal_and_df, python_type=StructuredDataset, expected=lt)
+        fdt.to_literal(ctx, sd_with_literal_and_df, python_type=StructuredDataset, expected=lt)
 
     sd_with_nothing = StructuredDataset()
     with pytest.raises(ValueError, match="If dataframe is not specified"):
-        FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_nothing, python_type=StructuredDataset, expected=lt)
+        fdt.to_literal(ctx, sd_with_nothing, python_type=StructuredDataset, expected=lt)
 
     sd_with_uri = StructuredDataset(uri="s3://some/extant/df.parquet")
 
     lt = TypeEngine.to_literal_type(Annotated[StructuredDataset, {}, "new-df-format"])
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt)
+    lit = fdt.to_literal(ctx, sd_with_uri, python_type=StructuredDataset, expected=lt)
     assert lit.scalar.structured_dataset.uri == "s3://some/extant/df.parquet"
     assert lit.scalar.structured_dataset.metadata.structured_dataset_type.format == "new-df-format"
 
@@ -190,8 +191,9 @@ def encode(
     assert lt.structured_dataset_type.format == "myavro"
 
     ctx = FlyteContextManager.current_context()
+    fdt = StructuredDatasetTransformerEngine()
     sd = StructuredDataset(dataframe=42)
-    l = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, sd, MyDF, lt)
+    l = fdt.to_literal(ctx, sd, MyDF, lt)
     # Test that the literal type is filled in even though the encode function above doesn't do it.
     assert l.scalar.structured_dataset.metadata.structured_dataset_type.format == "myavro"
 
@@ -272,24 +274,25 @@ def test_to_python_value_with_incoming_columns():
     ctx = FlyteContextManager.current_context()
     lt = TypeEngine.to_literal_type(original_type)
     df = generate_pandas()
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=original_type, expected=lt)
+    fdt = StructuredDatasetTransformerEngine()
+    lit = fdt.to_literal(ctx, df, python_type=original_type, expected=lt)
     assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 2
 
     # declare a new type that only has one column
     # get the dataframe, make sure it has the column that was asked for.
     subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)]
-    sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type)
+    sd = fdt.to_python_value(ctx, lit, subset_sd_type)
     assert sd.metadata.structured_dataset_type.columns[0].name == "age"
     sub_df = sd.open(pd.DataFrame).all()
     assert sub_df.shape[1] == 1
 
     # check when columns are not specified, should pull both and add column information.
-    sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset)
+    sd = fdt.to_python_value(ctx, lit, StructuredDataset)
     assert sd.metadata.structured_dataset_type.columns[0].name == "age"
 
     # should also work if subset type is just an annotated pd.DataFrame
     subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)]
-    sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type)
+    sub_df = fdt.to_python_value(ctx, lit, subset_pd_type)
     assert sub_df.shape[1] == 1
 
 
@@ -298,13 +301,14 @@ def test_to_python_value_without_incoming_columns():
     ctx = FlyteContextManager.current_context()
     lt = TypeEngine.to_literal_type(pd.DataFrame)
     df = generate_pandas()
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
+    fdt = StructuredDatasetTransformerEngine()
+    lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
     assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.columns) == 0
 
     # declare a new type that only has one column
     # get the dataframe, make sure it has the column that was asked for.
     subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)]
-    sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_sd_type)
+    sd = fdt.to_python_value(ctx, lit, subset_sd_type)
     assert sd.metadata.structured_dataset_type.columns[0].name == "age"
     sub_df = sd.open(pd.DataFrame).all()
     assert sub_df.shape[1] == 1
@@ -312,14 +316,14 @@ def test_to_python_value_without_incoming_columns():
     # check when columns are not specified, should pull both and add column information.
     # todo: see the todos in the open_as, and iter_as functions in StructuredDatasetTransformerEngine
     #  we have to recreate the literal because the test case above filled in the metadata
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
-    sd = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, StructuredDataset)
+    lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
+    sd = fdt.to_python_value(ctx, lit, StructuredDataset)
     assert sd.metadata.structured_dataset_type.columns == []
     sub_df = sd.open(pd.DataFrame).all()
     assert sub_df.shape[1] == 2
 
     # should also work if subset type is just an annotated pd.DataFrame
-    lit = FLYTE_DATASET_TRANSFORMER.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
+    lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
     subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)]
-    sub_df = FLYTE_DATASET_TRANSFORMER.to_python_value(ctx, lit, subset_pd_type)
+    sub_df = fdt.to_python_value(ctx, lit, subset_pd_type)
     assert sub_df.shape[1] == 1
diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
index 0a19f4fdf7..26152ff7e0 100644
--- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
+++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
@@ -18,7 +18,6 @@
 from flytekit.types.structured.structured_dataset import (
     BIGQUERY,
     DF,
-    FLYTE_DATASET_TRANSFORMER,
     StructuredDatasetTransformerEngine,
     LOCAL,
     PARQUET,

From 4ff5983bcfa743f2286cd6cd87ff079026fc6508 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Mon, 7 Feb 2022 18:07:43 -0800
Subject: [PATCH 07/11] revert warning to debug

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/types/structured/structured_dataset.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index c2059207aa..42c8076902 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -728,8 +728,8 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
 
 
 if USE_STRUCTURED_DATASET.get():
-    logger.warning("Structured dataset module load... using structured datasets!")
+    logger.debug("Structured dataset module load... using structured datasets!")
     flyte_dataset_transformer = StructuredDatasetTransformerEngine()
     TypeEngine.register(flyte_dataset_transformer)
 else:
-    logger.warning("Structured dataset module load... not using structured datasets")
+    logger.debug("Structured dataset module load... not using structured datasets")

From fa4a8097139a62f9da746cf94a0e4bf40841278a Mon Sep 17 00:00:00 2001
From: Kevin Su <pingsutw@apache.org>
Date: Tue, 8 Feb 2022 23:15:00 +0800
Subject: [PATCH 08/11] Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 flytekit/__init__.py                                |  2 +-
 flytekit/types/structured/structured_dataset.py     |  2 +-
 plugins/flytekit-papermill/dev-requirements.in      |  2 +-
 plugins/flytekit-papermill/dev-requirements.txt     | 13 ++++---------
 .../test_structured_dataset_workflow.py             |  2 +-
 5 files changed, 8 insertions(+), 13 deletions(-)

diff --git a/flytekit/__init__.py b/flytekit/__init__.py
index 303f34caa5..6e2e614702 100644
--- a/flytekit/__init__.py
+++ b/flytekit/__init__.py
@@ -193,8 +193,8 @@
     from flytekit.types.structured.structured_dataset import (
         StructuredDataset,
         StructuredDatasetFormat,
-        StructuredDatasetType,
         StructuredDatasetTransformerEngine,
+        StructuredDatasetType,
     )
 
 __version__ = "0.0.0+develop"
diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index 42c8076902..c4ca769c8c 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -24,7 +24,7 @@
 
 from flytekit.configuration.sdk import USE_STRUCTURED_DATASET
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
-from flytekit.core.type_engine import TypeTransformer, TypeEngine
+from flytekit.core.type_engine import TypeEngine, TypeTransformer
 from flytekit.loggers import logger
 from flytekit.models import literals
 from flytekit.models import types as type_models
diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in
index 8bf53e2636..209d7a9b9a 100644
--- a/plugins/flytekit-papermill/dev-requirements.in
+++ b/plugins/flytekit-papermill/dev-requirements.in
@@ -1 +1 @@
-flytekitplugins-spark>=0.30.0b4
+git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods
diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt
index 849cef4c9f..d80a261d94 100644
--- a/plugins/flytekit-papermill/dev-requirements.txt
+++ b/plugins/flytekit-papermill/dev-requirements.txt
@@ -40,9 +40,7 @@ docstring-parser==0.13
     # via flytekit
 flyteidl==0.21.23
     # via flytekit
-flytekit==0.26.0
-    # via flytekitplugins-spark
-flytekitplugins-spark==0.30.0b4
+flytekit @ git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods
     # via -r dev-requirements.in
 grpcio==1.43.0
     # via flytekit
@@ -87,12 +85,8 @@ protobuf==3.19.3
     #   flytekit
 py==1.11.0
     # via retry
-py4j==0.10.9.2
-    # via pyspark
 pyarrow==6.0.1
     # via flytekit
-pyspark==3.2.0
-    # via flytekitplugins-spark
 python-dateutil==2.8.1
     # via
     #   arrow
@@ -123,7 +117,6 @@ retry==0.9.2
 six==1.16.0
     # via
     #   cookiecutter
-    #   flytekit
     #   grpcio
     #   python-dateutil
     #   responses
@@ -134,7 +127,9 @@ statsd==3.3.0
 text-unidecode==1.3
     # via python-slugify
 typing-extensions==4.0.1
-    # via typing-inspect
+    # via
+    #   flytekit
+    #   typing-inspect
 typing-inspect==0.7.1
     # via dataclasses-json
 urllib3==1.26.8
diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
index 26152ff7e0..1d5006145e 100644
--- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
+++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
@@ -18,13 +18,13 @@
 from flytekit.types.structured.structured_dataset import (
     BIGQUERY,
     DF,
-    StructuredDatasetTransformerEngine,
     LOCAL,
     PARQUET,
     S3,
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
+    StructuredDatasetTransformerEngine,
 )
 
 PANDAS_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory()

From bff85be6460983f04bfb07397bff530f2fee2229 Mon Sep 17 00:00:00 2001
From: Kevin Su <pingsutw@apache.org>
Date: Tue, 8 Feb 2022 23:54:28 +0800
Subject: [PATCH 09/11] Fixed tests

Signed-off-by: Kevin Su <pingsutw@apache.org>
---
 plugins/flytekit-papermill/dev-requirements.in  | 4 +++-
 plugins/flytekit-papermill/dev-requirements.txt | 8 +++++++-
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in
index 209d7a9b9a..808bcd0f47 100644
--- a/plugins/flytekit-papermill/dev-requirements.in
+++ b/plugins/flytekit-papermill/dev-requirements.in
@@ -1 +1,3 @@
-git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods
+git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark
+# vcs+protocol://repo_url/#egg=pkg&subdirectory=flyte
+
diff --git a/plugins/flytekit-papermill/dev-requirements.txt b/plugins/flytekit-papermill/dev-requirements.txt
index d80a261d94..965cae8ad4 100644
--- a/plugins/flytekit-papermill/dev-requirements.txt
+++ b/plugins/flytekit-papermill/dev-requirements.txt
@@ -40,7 +40,9 @@ docstring-parser==0.13
     # via flytekit
 flyteidl==0.21.23
     # via flytekit
-flytekit @ git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods
+flytekit==0.30.0
+    # via flytekitplugins-spark
+flytekitplugins-spark @ git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#subdirectory=plugins/flytekit-spark
     # via -r dev-requirements.in
 grpcio==1.43.0
     # via flytekit
@@ -85,8 +87,12 @@ protobuf==3.19.3
     #   flytekit
 py==1.11.0
     # via retry
+py4j==0.10.9.3
+    # via pyspark
 pyarrow==6.0.1
     # via flytekit
+pyspark==3.2.1
+    # via flytekitplugins-spark
 python-dateutil==2.8.1
     # via
     #   arrow

From bd3be41121508ba2f82ae61647c1db8afc0400fe Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Tue, 8 Feb 2022 09:06:11 -0800
Subject: [PATCH 10/11] fix lint

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 plugins/flytekit-papermill/dev-requirements.in | 1 -
 1 file changed, 1 deletion(-)

diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in
index 808bcd0f47..32404e0a2d 100644
--- a/plugins/flytekit-papermill/dev-requirements.in
+++ b/plugins/flytekit-papermill/dev-requirements.in
@@ -1,3 +1,2 @@
 git+https://github.com/flyteorg/flytekit@add-sd-make-class-methods#egg=flytekitplugins-spark&subdirectory=plugins/flytekit-spark
 # vcs+protocol://repo_url/#egg=pkg&subdirectory=flyte
-

From d20109570bddc6ce2b70008f01b37d359ed7b2c4 Mon Sep 17 00:00:00 2001
From: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
Date: Tue, 8 Feb 2022 10:03:05 -0800
Subject: [PATCH 11/11] rename register_handler to register

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>
---
 flytekit/types/structured/basic_dfs.py             |  8 ++++----
 flytekit/types/structured/bigquery.py              |  8 ++++----
 flytekit/types/structured/structured_dataset.py    |  2 +-
 .../flytekitplugins/spark/sd_transformers.py       |  4 ++--
 .../flytekit/unit/core/test_structured_dataset.py  | 14 +++++++-------
 .../test_structured_dataset_workflow.py            |  8 ++++----
 6 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py
index 21bc08796a..49b2f13ed9 100644
--- a/flytekit/types/structured/basic_dfs.py
+++ b/flytekit/types/structured/basic_dfs.py
@@ -107,7 +107,7 @@ def decode(
 
 
 for protocol in [LOCAL, S3]:  # Should we add GCS
-    StructuredDatasetTransformerEngine.register_handler(PandasToParquetEncodingHandler(protocol), default_for_type=True)
-    StructuredDatasetTransformerEngine.register_handler(ParquetToPandasDecodingHandler(protocol), default_for_type=True)
-    StructuredDatasetTransformerEngine.register_handler(ArrowToParquetEncodingHandler(protocol), default_for_type=True)
-    StructuredDatasetTransformerEngine.register_handler(ParquetToArrowDecodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(protocol), default_for_type=True)
diff --git a/flytekit/types/structured/bigquery.py b/flytekit/types/structured/bigquery.py
index 8f1eb53925..923ea06e9e 100644
--- a/flytekit/types/structured/bigquery.py
+++ b/flytekit/types/structured/bigquery.py
@@ -110,7 +110,7 @@ def decode(
         return pa.Table.from_pandas(_read_from_bq(flyte_value))
 
 
-StructuredDatasetTransformerEngine.register_handler(PandasToBQEncodingHandlers(), default_for_type=False)
-StructuredDatasetTransformerEngine.register_handler(BQToPandasDecodingHandler(), default_for_type=False)
-StructuredDatasetTransformerEngine.register_handler(ArrowToBQEncodingHandlers(), default_for_type=False)
-StructuredDatasetTransformerEngine.register_handler(BQToArrowDecodingHandler(), default_for_type=False)
+StructuredDatasetTransformerEngine.register(PandasToBQEncodingHandlers(), default_for_type=False)
+StructuredDatasetTransformerEngine.register(BQToPandasDecodingHandler(), default_for_type=False)
+StructuredDatasetTransformerEngine.register(ArrowToBQEncodingHandlers(), default_for_type=False)
+StructuredDatasetTransformerEngine.register(BQToArrowDecodingHandler(), default_for_type=False)
diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py
index c4ca769c8c..819bc012cc 100644
--- a/flytekit/types/structured/structured_dataset.py
+++ b/flytekit/types/structured/structured_dataset.py
@@ -380,7 +380,7 @@ def __init__(self):
         self._type_assertions_enabled = False
 
     @classmethod
-    def register_handler(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False):
+    def register(cls, h: Handlers, default_for_type: Optional[bool] = True, override: Optional[bool] = False):
         """
         Call this with any handler to register it with this dataframe meta-transformer
 
diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
index 28e2ffdadc..2466d3fc13 100644
--- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
+++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py
@@ -45,5 +45,5 @@ def decode(
 
 
 for protocol in ["/", "s3"]:
-    StructuredDatasetTransformerEngine.register_handler(SparkToParquetEncodingHandler(protocol), default_for_type=True)
-    StructuredDatasetTransformerEngine.register_handler(ParquetToSparkDecodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(SparkToParquetEncodingHandler(protocol), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(ParquetToSparkDecodingHandler(protocol), default_for_type=True)
diff --git a/tests/flytekit/unit/core/test_structured_dataset.py b/tests/flytekit/unit/core/test_structured_dataset.py
index 3d66e01b29..a7ef1ea953 100644
--- a/tests/flytekit/unit/core/test_structured_dataset.py
+++ b/tests/flytekit/unit/core/test_structured_dataset.py
@@ -129,15 +129,15 @@ def __init__(self, protocol):
         def encode(self):
             ...
 
-    StructuredDatasetTransformerEngine.register_handler(TempEncoder("gs"), default_for_type=False)
+    StructuredDatasetTransformerEngine.register(TempEncoder("gs"), default_for_type=False)
     with pytest.raises(ValueError):
-        StructuredDatasetTransformerEngine.register_handler(TempEncoder("gs://"), default_for_type=False)
+        StructuredDatasetTransformerEngine.register(TempEncoder("gs://"), default_for_type=False)
 
     class TempEncoder:
         pass
 
     with pytest.raises(TypeError, match="We don't support this type of handler"):
-        StructuredDatasetTransformerEngine.register_handler(TempEncoder, default_for_type=False)
+        StructuredDatasetTransformerEngine.register(TempEncoder, default_for_type=False)
 
 
 def test_to_literal():
@@ -186,7 +186,7 @@ def encode(
         ) -> literals.StructuredDataset:
             return literals.StructuredDataset(uri="")
 
-    StructuredDatasetTransformerEngine.register_handler(TempEncoder("myavro"), default_for_type=True)
+    StructuredDatasetTransformerEngine.register(TempEncoder("myavro"), default_for_type=True)
     lt = TypeEngine.to_literal_type(MyDF)
     assert lt.structured_dataset_type.format == "myavro"
 
@@ -199,7 +199,7 @@ def encode(
 
     # Test that looking up encoders/decoders falls back to the "" encoder/decoder
     empty_format_temp_encoder = TempEncoder("")
-    StructuredDatasetTransformerEngine.register_handler(empty_format_temp_encoder, default_for_type=False)
+    StructuredDatasetTransformerEngine.register(empty_format_temp_encoder, default_for_type=False)
 
     res = StructuredDatasetTransformerEngine.get_encoder(MyDF, "tmpfs", "rando")
     assert res is empty_format_temp_encoder
@@ -224,7 +224,7 @@ def decode(
         ) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]:
             yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
 
-    StructuredDatasetTransformerEngine.register_handler(
+    StructuredDatasetTransformerEngine.register(
         MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False
     )
     sd = StructuredDataset()
@@ -244,7 +244,7 @@ def decode(
         ) -> pd.DataFrame:
             pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
 
-    StructuredDatasetTransformerEngine.register_handler(
+    StructuredDatasetTransformerEngine.register(
         MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True
     )
     sd = StructuredDataset()
diff --git a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
index 1d5006145e..d911f971e4 100644
--- a/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
+++ b/tests/flytekit/unit/types/structured_dataset/test_structured_dataset_workflow.py
@@ -58,8 +58,8 @@ def decode(
         return pd_df
 
 
-StructuredDatasetTransformerEngine.register_handler(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True)
-StructuredDatasetTransformerEngine.register_handler(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True)
+StructuredDatasetTransformerEngine.register(MockBQEncodingHandlers(pd.DataFrame, BIGQUERY), False, True)
+StructuredDatasetTransformerEngine.register(MockBQDecodingHandlers(pd.DataFrame, BIGQUERY), False, True)
 
 
 class NumpyEncodingHandlers(StructuredDatasetEncoder):
@@ -95,8 +95,8 @@ def decode(
 
 
 for protocol in [LOCAL, S3]:
-    StructuredDatasetTransformerEngine.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET))
-    StructuredDatasetTransformerEngine.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET))
+    StructuredDatasetTransformerEngine.register(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET))
+    StructuredDatasetTransformerEngine.register(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET))
 
 
 @task
