From a03bc8f98c8ed740ff0f8ed6f6879a7d3a05f88c Mon Sep 17 00:00:00 2001 From: Yee Hing Tong Date: Mon, 7 Feb 2022 10:40:38 -0800 Subject: [PATCH 01/11] import singleton Signed-off-by: Yee Hing Tong --- 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 Date: Mon, 7 Feb 2022 14:15:52 -0800 Subject: [PATCH 02/11] make compat a different job Signed-off-by: Yee Hing Tong --- .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 Date: Mon, 7 Feb 2022 14:39:21 -0800 Subject: [PATCH 03/11] nit Signed-off-by: Yee Hing Tong --- 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 Date: Mon, 7 Feb 2022 15:57:04 -0800 Subject: [PATCH 04/11] wrong import location Signed-off-by: Yee Hing Tong --- 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 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 --- 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 Date: Mon, 7 Feb 2022 18:02:59 -0800 Subject: [PATCH 06/11] make class Signed-off-by: Yee Hing Tong --- 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 Date: Mon, 7 Feb 2022 18:07:43 -0800 Subject: [PATCH 07/11] revert warning to debug Signed-off-by: Yee Hing Tong --- 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 Date: Tue, 8 Feb 2022 23:15:00 +0800 Subject: [PATCH 08/11] Fixed tests Signed-off-by: Kevin Su --- 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 Date: Tue, 8 Feb 2022 23:54:28 +0800 Subject: [PATCH 09/11] Fixed tests Signed-off-by: Kevin Su --- 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 Date: Tue, 8 Feb 2022 09:06:11 -0800 Subject: [PATCH 10/11] fix lint Signed-off-by: Yee Hing Tong --- 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 Date: Tue, 8 Feb 2022 10:03:05 -0800 Subject: [PATCH 11/11] rename register_handler to register Signed-off-by: Yee Hing Tong --- 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