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 diff --git a/flytekit/__init__.py b/flytekit/__init__.py index 64c0a26530..6e2e614702 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -193,6 +193,7 @@ from flytekit.types.structured.structured_dataset import ( StructuredDataset, StructuredDatasetFormat, + StructuredDatasetTransformerEngine, StructuredDatasetType, ) diff --git a/flytekit/types/structured/basic_dfs.py b/flytekit/types/structured/basic_dfs.py index a65fbbeee3..49b2f13ed9 100644 --- a/flytekit/types/structured/basic_dfs.py +++ b/flytekit/types/structured/basic_dfs.py @@ -12,13 +12,13 @@ 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, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) T = TypeVar("T") @@ -107,7 +107,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(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 d9221e4110..923ea06e9e 100644 --- a/flytekit/types/structured/bigquery.py +++ b/flytekit/types/structured/bigquery.py @@ -12,11 +12,11 @@ from flytekit.types.structured.structured_dataset import ( BIGQUERY, DF, - FLYTE_DATASET_TRANSFORMER, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, StructuredDatasetMetadata, + StructuredDatasetTransformerEngine, ) @@ -110,7 +110,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(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 ffefd101dd..819bc012cc 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 TypeEngine, TypeTransformer from flytekit.loggers import logger from flytekit.models import literals from flytekit.models import types as type_models @@ -106,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 ) @@ -114,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 ) @@ -170,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 @@ -230,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 @@ -337,7 +336,8 @@ class StructuredDatasetTransformerEngine(TypeTransformer[StructuredDataset]): 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: @@ -352,18 +352,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: @@ -376,7 +379,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(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 @@ -386,7 +390,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 @@ -394,13 +398,14 @@ 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) + engine = StructuredDatasetTransformerEngine() + TypeEngine.register_additional_type(engine, h.python_type, override=True) def assert_type(self, t: Type[StructuredDataset], v: typing.Any): return @@ -724,7 +729,7 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if USE_STRUCTURED_DATASET.get(): logger.debug("Structured dataset module load... using structured datasets!") - FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() - TypeEngine.register(FLYTE_DATASET_TRANSFORMER) + flyte_dataset_transformer = StructuredDatasetTransformerEngine() + TypeEngine.register(flyte_dataset_transformer) else: logger.debug("Structured dataset module load... not using structured datasets") diff --git a/plugins/flytekit-papermill/dev-requirements.in b/plugins/flytekit-papermill/dev-requirements.in index 8bf53e2636..32404e0a2d 100644 --- a/plugins/flytekit-papermill/dev-requirements.in +++ b/plugins/flytekit-papermill/dev-requirements.in @@ -1 +1,2 @@ -flytekitplugins-spark>=0.30.0b4 +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 849cef4c9f..965cae8ad4 100644 --- a/plugins/flytekit-papermill/dev-requirements.txt +++ b/plugins/flytekit-papermill/dev-requirements.txt @@ -40,9 +40,9 @@ docstring-parser==0.13 # via flytekit flyteidl==0.21.23 # via flytekit -flytekit==0.26.0 +flytekit==0.30.0 # via flytekitplugins-spark -flytekitplugins-spark==0.30.0b4 +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 @@ -87,11 +87,11 @@ protobuf==3.19.3 # flytekit py==1.11.0 # via retry -py4j==0.10.9.2 +py4j==0.10.9.3 # via pyspark pyarrow==6.0.1 # via flytekit -pyspark==3.2.0 +pyspark==3.2.1 # via flytekitplugins-spark python-dateutil==2.8.1 # via @@ -123,7 +123,6 @@ retry==0.9.2 six==1.16.0 # via # cookiecutter - # flytekit # grpcio # python-dateutil # responses @@ -134,7 +133,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/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py index e0b1c7b41e..2466d3fc13 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -7,11 +7,11 @@ 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, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) @@ -45,5 +45,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(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 4e4309e292..a7ef1ea953 100644 --- a/tests/flytekit/unit/core/test_structured_dataset.py +++ b/tests/flytekit/unit/core/test_structured_dataset.py @@ -19,11 +19,11 @@ from flytekit import kwtypes from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, PARQUET, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, convert_schema_type_to_structured_dataset_type, extract_cols_and_format, protocol_prefix, @@ -117,10 +117,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 +129,15 @@ def __init__(self, protocol): def encode(self): ... - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder("gs"), default_for_type=False) + StructuredDatasetTransformerEngine.register(TempEncoder("gs"), default_for_type=False) with pytest.raises(ValueError): - FLYTE_DATASET_TRANSFORMER.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"): - FLYTE_DATASET_TRANSFORMER.register_handler(TempEncoder, default_for_type=False) + StructuredDatasetTransformerEngine.register(TempEncoder, default_for_type=False) def test_to_literal(): @@ -145,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 @@ -153,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" @@ -184,21 +186,22 @@ def encode( ) -> literals.StructuredDataset: return literals.StructuredDataset(uri="") - FLYTE_DATASET_TRANSFORMER.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" 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" # 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(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 +224,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( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False ) sd = StructuredDataset() @@ -241,7 +244,7 @@ def decode( ) -> pd.DataFrame: pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]}) - FLYTE_DATASET_TRANSFORMER.register_handler( + StructuredDatasetTransformerEngine.register( MockPandasDecodingHandlers(pd.DataFrame, "tmpfs"), default_for_type=False, override=True ) sd = StructuredDataset() @@ -271,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 @@ -297,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 @@ -311,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 fa879ff0d9..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 @@ -18,13 +18,13 @@ from flytekit.types.structured.structured_dataset import ( BIGQUERY, DF, - FLYTE_DATASET_TRANSFORMER, LOCAL, PARQUET, S3, StructuredDataset, StructuredDatasetDecoder, StructuredDatasetEncoder, + StructuredDatasetTransformerEngine, ) PANDAS_PATH = FlyteContextManager.current_context().file_access.get_random_local_directory() @@ -58,8 +58,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(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]: - FLYTE_DATASET_TRANSFORMER.register_handler(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET)) - FLYTE_DATASET_TRANSFORMER.register_handler(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET)) + StructuredDatasetTransformerEngine.register(NumpyEncodingHandlers(np.ndarray, protocol, PARQUET)) + StructuredDatasetTransformerEngine.register(NumpyDecodingHandlers(np.ndarray, protocol, PARQUET)) @task