diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml index a8f4c06b9a..bedf66383e 100644 --- a/.github/workflows/pythonbuild.yml +++ b/.github/workflows/pythonbuild.yml @@ -41,13 +41,14 @@ jobs: pip freeze - name: Test with coverage run: | - coverage run -m pytest tests/flytekit/unit + 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 shell: 'script -q -e -c "bash {0}"' run: | python -m pip install awscli - coverage run --append -m pytest tests/flytekit/integration + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run --append -m pytest tests/flytekit/integration - name: Codecov uses: codecov/codecov-action@v1 with: @@ -108,7 +109,7 @@ jobs: - name: Test with coverage run: | cd plugins/${{ matrix.plugin-names }} - coverage run -m pytest tests + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests lint: runs-on: ubuntu-latest diff --git a/Dockerfile.py310 b/Dockerfile.py310 index 1dc6b117b0..8fc7ae9e3d 100644 --- a/Dockerfile.py310 +++ b/Dockerfile.py310 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py37 b/Dockerfile.py37 index 4eb531b8de..ae59a7d608 100644 --- a/Dockerfile.py37 +++ b/Dockerfile.py37 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py38 b/Dockerfile.py38 index 934f167748..bdc9acec05 100644 --- a/Dockerfile.py38 +++ b/Dockerfile.py38 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Dockerfile.py39 b/Dockerfile.py39 index 1a4d617aa0..6e972b2589 100644 --- a/Dockerfile.py39 +++ b/Dockerfile.py39 @@ -14,3 +14,4 @@ RUN pip install -U flytekit==$VERSION WORKDIR /app ENV FLYTE_INTERNAL_IMAGE "$DOCKER_IMAGE" +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE diff --git a/Makefile b/Makefile index b1088aece7..10f4e5b232 100644 --- a/Makefile +++ b/Makefile @@ -46,12 +46,11 @@ spellcheck: ## Runs a spellchecker over all code and documentation codespell -L "te,raison,fo" --skip="./docs/build,./.git" .PHONY: test -test: lint ## Run tests - pytest tests/flytekit/unit +test: lint unit_test .PHONY: unit_test unit_test: - pytest tests/flytekit/unit + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE pytest tests/flytekit/unit tests/flytekit_compatibility requirements-spark2.txt: export CUSTOM_COMPILE_COMMAND := make requirements-spark2.txt requirements-spark2.txt: requirements-spark2.in install-piptools @@ -79,7 +78,7 @@ requirements: requirements.txt dev-requirements.txt requirements-spark2.txt doc- # TODO: Change this in the future to be all of flytekit .PHONY: coverage coverage: - coverage run -m pytest tests/flytekit/unit/core flytekit/types + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE coverage run -m pytest tests/flytekit/unit/core flytekit/types coverage report -m --include="flytekit/core/*,flytekit/types/*" PLACEHOLDER := "__version__\ =\ \"0.0.0+develop\"" diff --git a/flytekit/__init__.py b/flytekit/__init__.py index fd8fcc7d9c..dcf8878533 100644 --- a/flytekit/__init__.py +++ b/flytekit/__init__.py @@ -160,6 +160,7 @@ else: from importlib.metadata import entry_points +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.core.base_sql_task import SQLTask from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes from flytekit.core.condition import conditional @@ -186,11 +187,13 @@ from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar from flytekit.models.types import LiteralType from flytekit.types import directory, file, schema -from flytekit.types.structured.structured_dataset import ( - StructuredDataset, - StructuredDatasetFormat, - StructuredDatasetType, -) + +if USE_STRUCTURED_DATASET.get(): + from flytekit.types.structured.structured_dataset import ( + StructuredDataset, + StructuredDatasetFormat, + StructuredDatasetType, + ) __version__ = "0.0.0+develop" diff --git a/flytekit/configuration/platform.py b/flytekit/configuration/platform.py index 5c4061fa4f..eecbeda162 100644 --- a/flytekit/configuration/platform.py +++ b/flytekit/configuration/platform.py @@ -1,5 +1,4 @@ from flytekit.configuration import common as _config_common -from flytekit.core import constants as _constants URL = _config_common.FlyteStringConfigurationEntry("platform", "url") @@ -14,10 +13,6 @@ INSECURE = _config_common.FlyteBoolConfigurationEntry("platform", "insecure", default=False) -CLOUD_PROVIDER = _config_common.FlyteStringConfigurationEntry( - "platform", "cloud_provider", default=_constants.CloudProvider.AWS -) - AUTH = _config_common.FlyteBoolConfigurationEntry("platform", "auth", default=False) """ This config setting should not normally be filled in. Whether or not an admin server requires authentication should be diff --git a/flytekit/configuration/sdk.py b/flytekit/configuration/sdk.py index 7e5c8872ce..a142b29546 100644 --- a/flytekit/configuration/sdk.py +++ b/flytekit/configuration/sdk.py @@ -8,20 +8,6 @@ and execution of entities. """ -EXECUTION_ENGINE = _config_common.FlyteStringConfigurationEntry("sdk", "execution_engine", default="flyte") -""" -This is a comma-delimited list of package strings, in order, for resolving execution behavior. - -TODO: Explain how this would be used to extend the SDK -""" - -TYPE_ENGINES = _config_common.FlyteStringListConfigurationEntry("sdk", "type_engines", default=[]) -""" -This is a comma-delimited list of package strings, in order, for resolving type behavior. - -TODO: Explain how this would be used to extend the SDK -""" - LOCAL_SANDBOX = _config_common.FlyteStringConfigurationEntry( "sdk", "local_sandbox", @@ -32,47 +18,6 @@ clean up data in these directories. """ -SDK_PYTHON_VENV = _config_common.FlyteStringListConfigurationEntry("sdk", "python_venv", default=[]) -""" -This is a list of commands/args which will be prefixed to the entrypoint command by SDK. -""" - -ROLE = _config_common.FlyteStringConfigurationEntry("sdk", "role") -""" -This is the role the SDK will use by default to execute workflows. For example, in AWS this should be an IAM role -string. -""" - -NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "name_format", default="{module}.{name}") -""" -This is a Python format string which the SDK will use to generate names for discovered entities. The default is -'{module}.{name}' which will result in strings like 'package.module.name'. Any template portion of the string can only -include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. -""" - -TASK_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "task_name_format", fallback=NAME_FORMAT) -""" -This is a Python format string which the SDK will use to generate names for tasks. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - -WORKFLOW_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry("sdk", "workflow_name_format", fallback=NAME_FORMAT) -""" -This is a Python format string which the SDK will use to generate names for workflows. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - -LAUNCH_PLAN_NAME_FORMAT = _config_common.FlyteStringConfigurationEntry( - "sdk", "launch_plan_name_format", fallback=NAME_FORMAT -) -""" -This is a Python format string which the SDK will use to generate names for launch plans. Any template portion of the -string can only include 'module' or 'name'. So '{name}' is valid, but '{key}' is not. If not specified, -we fall back to the configuration for :py:attr:`flytekit.configuration.sdk.NAME_FORMAT` -""" - LOGGING_LEVEL = _config_common.FlyteIntegerConfigurationEntry("sdk", "logging_level", default=20) """ This is the default logging level for the Python logging library and will be set before user code runs. @@ -85,9 +30,9 @@ This is the parquet engine to use when reading data from parquet files. """ -FAST_REGISTRATION_DIR = _config_common.FlyteStringConfigurationEntry("sdk", "fast_registration_dir") +# Feature Gate +USE_STRUCTURED_DATASET = _config_common.FlyteBoolConfigurationEntry("sdk", "use_structured_dataset", default=False) """ -This is the remote directory where fast-registered code will be uploaded to. -Users calling fast-execute need write permission to this directory. -Furthermore, it is important that whichever role executes your workflow has read access to this directory. +Note: This gate will be switched to True at some point in the future. Definitely by 1.0, if not v0.31.0. + """ diff --git a/flytekit/types/structured/__init__.py b/flytekit/types/structured/__init__.py index da80c89016..c4fd015c57 100644 --- a/flytekit/types/structured/__init__.py +++ b/flytekit/types/structured/__init__.py @@ -1,21 +1,23 @@ +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET from flytekit.loggers import logger -from .basic_dfs import ( - ArrowToParquetEncodingHandler, - PandasToParquetEncodingHandler, - ParquetToArrowDecodingHandler, - ParquetToPandasDecodingHandler, -) - -try: - from .bigquery import ( - ArrowToBQEncodingHandlers, - BQToArrowDecodingHandler, - BQToPandasDecodingHandler, - PandasToBQEncodingHandlers, - ) -except ImportError: - logger.info( - "We won't register bigquery handler for structured dataset because " - "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery" +if USE_STRUCTURED_DATASET.get(): + from .basic_dfs import ( + ArrowToParquetEncodingHandler, + PandasToParquetEncodingHandler, + ParquetToArrowDecodingHandler, + ParquetToPandasDecodingHandler, ) + + try: + from .bigquery import ( + ArrowToBQEncodingHandlers, + BQToArrowDecodingHandler, + BQToPandasDecodingHandler, + PandasToBQEncodingHandlers, + ) + except ImportError: + logger.info( + "We won't register bigquery handler for structured dataset because " + "we can't find the packages google-cloud-bigquery-storage and google-cloud-bigquery" + ) diff --git a/flytekit/types/structured/structured_dataset.py b/flytekit/types/structured/structured_dataset.py index a36e50b976..ffefd101dd 100644 --- a/flytekit/types/structured/structured_dataset.py +++ b/flytekit/types/structured/structured_dataset.py @@ -22,6 +22,7 @@ import numpy as _np import pyarrow as pa +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 @@ -381,6 +382,10 @@ def register_handler(self, h: Handlers, default_for_type: Optional[bool] = True, The string "://" should not be present in any handler's protocol so we don't check for it. """ + if not USE_STRUCTURED_DATASET.get(): + logger.info(f"Structured datasets not enabled, not registering handler {h}") + return + lowest_level = self._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)}") @@ -717,5 +722,9 @@ def guess_python_type(self, literal_type: LiteralType) -> Type[T]: raise ValueError(f"StructuredDatasetTransformerEngine cannot reverse {literal_type}") -FLYTE_DATASET_TRANSFORMER = StructuredDatasetTransformerEngine() -TypeEngine.register(FLYTE_DATASET_TRANSFORMER) +if USE_STRUCTURED_DATASET.get(): + logger.debug("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") diff --git a/plugins/Makefile b/plugins/Makefile index 9131836c77..94dabbde85 100644 --- a/plugins/Makefile +++ b/plugins/Makefile @@ -1,6 +1,6 @@ .PHONY: test test: - find . -maxdepth 1 -type d | grep 'flytekit-' | xargs -L1 pytest + FLYTE_SDK_USE_STRUCTURED_DATASET=TRUE find . -maxdepth 1 -type d | grep 'flytekit-' | xargs -L1 pytest .PHONY: build_all_plugins build_all_plugins: diff --git a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py index 145497b030..d239632248 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/__init__.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/__init__.py @@ -1,2 +1,7 @@ -from .schema import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler +from flytekit.configuration.sdk import USE_STRUCTURED_DATASET + +from .schema import SparkDataFrameSchemaReader, SparkDataFrameSchemaWriter, SparkDataFrameTransformer # noqa from .task import Spark, new_spark_session + +if USE_STRUCTURED_DATASET.get(): + from .sd_transformers import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler diff --git a/plugins/flytekit-spark/flytekitplugins/spark/schema.py b/plugins/flytekit-spark/flytekitplugins/spark/schema.py index 72dad4fc92..1cae101295 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/schema.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/schema.py @@ -2,21 +2,12 @@ from typing import Type import pyspark -from pyspark.sql.dataframe import DataFrame from flytekit import FlyteContext from flytekit.extend import T, TypeEngine, TypeTransformer -from flytekit.models import literals -from flytekit.models.literals import Literal, Scalar, Schema, StructuredDatasetMetadata -from flytekit.models.types import LiteralType, SchemaType, StructuredDatasetType +from flytekit.models.literals import Literal, Scalar, Schema +from flytekit.models.types import LiteralType, SchemaType from flytekit.types.schema import SchemaEngine, SchemaFormat, SchemaHandler, SchemaReader, SchemaWriter -from flytekit.types.structured.structured_dataset import ( - FLYTE_DATASET_TRANSFORMER, - PARQUET, - StructuredDataset, - StructuredDatasetDecoder, - StructuredDatasetEncoder, -) class SparkDataFrameSchemaReader(SchemaReader[pyspark.sql.DataFrame]): @@ -106,37 +97,3 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: # %% # This makes pyspark.DataFrame as a supported output/input type with flytekit. TypeEngine.register(SparkDataFrameTransformer()) - - -class SparkToParquetEncodingHandler(StructuredDatasetEncoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def encode( - self, - ctx: FlyteContext, - structured_dataset: StructuredDataset, - structured_dataset_type: StructuredDatasetType, - ) -> literals.StructuredDataset: - path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() - df = typing.cast(DataFrame, structured_dataset.dataframe) - df.write.mode("overwrite").parquet(path) - return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) - - -class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): - def __init__(self, protocol: str): - super().__init__(DataFrame, protocol, PARQUET) - - def decode( - self, - ctx: FlyteContext, - flyte_value: literals.StructuredDataset, - ) -> DataFrame: - user_ctx = FlyteContext.current_context().user_space_params - return user_ctx.spark_session.read.parquet(flyte_value.uri) - - -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) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py new file mode 100644 index 0000000000..e0b1c7b41e --- /dev/null +++ b/plugins/flytekit-spark/flytekitplugins/spark/sd_transformers.py @@ -0,0 +1,49 @@ +import typing + +from pyspark.sql.dataframe import DataFrame + +from flytekit import FlyteContext +from flytekit.models import literals +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, +) + + +class SparkToParquetEncodingHandler(StructuredDatasetEncoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def encode( + self, + ctx: FlyteContext, + structured_dataset: StructuredDataset, + structured_dataset_type: StructuredDatasetType, + ) -> literals.StructuredDataset: + path = typing.cast(str, structured_dataset.uri) or ctx.file_access.get_random_remote_directory() + df = typing.cast(DataFrame, structured_dataset.dataframe) + df.write.mode("overwrite").parquet(path) + return literals.StructuredDataset(uri=path, metadata=StructuredDatasetMetadata(structured_dataset_type)) + + +class ParquetToSparkDecodingHandler(StructuredDatasetDecoder): + def __init__(self, protocol: str): + super().__init__(DataFrame, protocol, PARQUET) + + def decode( + self, + ctx: FlyteContext, + flyte_value: literals.StructuredDataset, + ) -> DataFrame: + user_ctx = FlyteContext.current_context().user_space_params + return user_ctx.spark_session.read.parquet(flyte_value.uri) + + +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) diff --git a/plugins/flytekit-sqlalchemy/Dockerfile b/plugins/flytekit-sqlalchemy/Dockerfile index b1c47ffd1e..74c9678746 100644 --- a/plugins/flytekit-sqlalchemy/Dockerfile +++ b/plugins/flytekit-sqlalchemy/Dockerfile @@ -5,6 +5,7 @@ ENV VENV /opt/venv ENV LANG C.UTF-8 ENV LC_ALL C.UTF-8 ENV PYTHONPATH /app +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE RUN pip install awscli RUN pip install gsutil diff --git a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile index 7e5d01829f..5548412616 100644 --- a/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile +++ b/tests/flytekit/integration/remote/mock_flyte_repo/workflows/Dockerfile @@ -6,6 +6,7 @@ ENV VENV /opt/venv ENV LANG C.UTF-8 ENV LC_ALL C.UTF-8 ENV PYTHONPATH /root +ENV FLYTE_SDK_USE_STRUCTURED_DATASET TRUE # This is necessary for opencv to work RUN apt-get update && apt-get install -y libsm6 libxext6 libxrender-dev ffmpeg build-essential diff --git a/tests/flytekit_compatibility/__init__.py b/tests/flytekit_compatibility/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit_compatibility/test_schema_types.py b/tests/flytekit_compatibility/test_schema_types.py new file mode 100644 index 0000000000..6711ec34de --- /dev/null +++ b/tests/flytekit_compatibility/test_schema_types.py @@ -0,0 +1,53 @@ +from datetime import datetime, timedelta + +import pytest + +from flytekit import kwtypes +from flytekit.core import context_manager +from flytekit.core.context_manager import ExecutionState +from flytekit.core.type_engine import TypeEngine +from flytekit.types.schema import FlyteSchema, SchemaFormat +from flytekit.types.schema.types import FlyteSchemaTransformer + + +def test_typed_schema(): + s = FlyteSchema[kwtypes(x=int, y=float)] + assert s.format() == SchemaFormat.PARQUET + assert s.columns() == {"x": int, "y": float} + + +def test_assert_type(): + ctx = context_manager.FlyteContextManager.current_context() + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state(ctx.new_execution_state().with_params(mode=ExecutionState.Mode.TASK_EXECUTION)) + ) as ctx: + schema = FlyteSchema[kwtypes(x=int, y=float)] + fst = FlyteSchemaTransformer() + lt = fst.get_literal_type(schema) + with pytest.raises(ValueError, match="DataFrames of type are not supported currently"): + TypeEngine.to_literal(ctx, 3, schema, lt) + + +def test_schema_back_and_forth(): + orig = FlyteSchema[kwtypes(TrackId=int, Name=str)] + lt = TypeEngine.to_literal_type(orig) + pt = TypeEngine.guess_python_type(lt) + lt2 = TypeEngine.to_literal_type(pt) + assert lt == lt2 + + +def test_remaining_prims(): + orig = FlyteSchema[kwtypes(my_dt=datetime, my_td=timedelta, my_b=bool)] + lt = TypeEngine.to_literal_type(orig) + pt = TypeEngine.guess_python_type(lt) + lt2 = TypeEngine.to_literal_type(pt) + assert lt == lt2 + + +def test_bad_conversion(): + orig = FlyteSchema[kwtypes(my_custom=bool)] + lt = TypeEngine.to_literal_type(orig) + # Make a not real column type + lt.schema.columns[0]._type = 15 + with pytest.raises(ValueError): + TypeEngine.guess_python_type(lt) diff --git a/tests/flytekit_compatibility/test_schema_usage_copied.py b/tests/flytekit_compatibility/test_schema_usage_copied.py new file mode 100644 index 0000000000..dfeee49cf3 --- /dev/null +++ b/tests/flytekit_compatibility/test_schema_usage_copied.py @@ -0,0 +1,149 @@ +import datetime +from dataclasses import dataclass + +import pandas as pd +from dataclasses_json import dataclass_json + +from flytekit import SQLTask, kwtypes +from flytekit.core import context_manager +from flytekit.core.task import TaskMetadata, task +from flytekit.core.testing import patch, task_mock +from flytekit.core.workflow import workflow +from flytekit.types.schema import FlyteSchema, SchemaOpenMode + + +def test_wf1_with_sql(): + sql = SQLTask( + "my-query", + query_template="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds }}' LIMIT 10", + inputs=kwtypes(ds=datetime.datetime), + outputs=kwtypes(results=FlyteSchema), + metadata=TaskMetadata(retries=2), + ) + + @task + def t1() -> datetime.datetime: + return datetime.datetime.now() + + @workflow + def my_wf() -> FlyteSchema: + dt = t1() + return sql(ds=dt) + + with task_mock(sql) as mock: + mock.return_value = pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}) + assert (my_wf().open().all() == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})).all().all() + assert context_manager.FlyteContextManager.size() == 1 + + +def test_wf1_with_sql_with_patch(): + sql = SQLTask( + "my-query", + query_template="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds }}' LIMIT 10", + inputs=kwtypes(ds=datetime.datetime), + outputs=kwtypes(results=FlyteSchema), + metadata=TaskMetadata(retries=2), + ) + + @task + def t1() -> datetime.datetime: + return datetime.datetime.now() + + @workflow + def my_wf() -> FlyteSchema: + dt = t1() + return sql(ds=dt) + + @patch(sql) + def test_user_demo_test(mock_sql): + mock_sql.return_value = pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}) + assert (my_wf().open().all() == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})).all().all() + + # Have to call because tests inside tests don't run + test_user_demo_test() + assert context_manager.FlyteContextManager.size() == 1 + + +def test_wf_typed_schema(): + schema1 = FlyteSchema[kwtypes(x=int, y=str)] + + @task + def t1() -> schema1: + s = schema1() + s.open().write(pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) + return s + + @task + def t2(s: FlyteSchema[kwtypes(x=int, y=str)]) -> FlyteSchema[kwtypes(x=int)]: + df = s.open().all() + return df[s.column_names()[:-1]] + + @workflow + def wf() -> FlyteSchema[kwtypes(x=int)]: + return t2(s=t1()) + + w = t1() + assert w is not None + df = w.open(override_mode=SchemaOpenMode.READ).all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]}).reset_index(drop=True) + assert result_df.all().all() + + df = t2(s=w.as_readonly()) + df = df.open(override_mode=SchemaOpenMode.READ).all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2]}).reset_index(drop=True) + assert result_df.all().all() + + x = wf() + df = x.open().all() + result_df = df.reset_index(drop=True) == pd.DataFrame(data={"x": [1, 2]}).reset_index(drop=True) + assert result_df.all().all() + + +def test_wf_schema_to_df(): + schema1 = FlyteSchema[kwtypes(x=int, y=str)] + + @task + def t1() -> schema1: + s = schema1() + s.open().write(pd.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})) + return s + + @task + def t2(df: pd.DataFrame) -> int: + return len(df.columns.values) + + @workflow + def wf() -> int: + return t2(df=t1()) + + x = wf() + assert x == 2 + + +def test_flyte_schema_dataclass(): + TestSchema = FlyteSchema[kwtypes(some_str=str)] + + @dataclass_json + @dataclass + class InnerResult: + number: int + schema: TestSchema + + @dataclass_json + @dataclass + class Result: + result: InnerResult + schema: TestSchema + + schema = TestSchema() + + @task + def t1(x: int) -> Result: + + return Result(result=InnerResult(number=x, schema=schema), schema=schema) + + @workflow + def wf(x: int) -> Result: + return t1(x=x) + + assert wf(x=10) == Result(result=InnerResult(number=10, schema=schema), schema=schema) diff --git a/tests/flytekit_compatibility/test_structured_dataset.py b/tests/flytekit_compatibility/test_structured_dataset.py new file mode 100644 index 0000000000..20965cb802 --- /dev/null +++ b/tests/flytekit_compatibility/test_structured_dataset.py @@ -0,0 +1,8 @@ +import pandas as pd + +from flytekit.core.type_engine import TypeEngine + + +def test_pandas_is_schema_with_flag(): + lt = TypeEngine.to_literal_type(pd.DataFrame) + assert lt.schema is not None