def test_nonfunction_task_and_df_input(): @reference_task( project="flytesnacks", domain="development", name="ref_t1", version="fast56d8ce2e373baf011f4d3532e45f0a9b", ) def ref_t1( dataframe: pd.DataFrame, imputation_method: str = "median", ) -> pd.DataFrame: ... @reference_task( project="flytesnacks", domain="development", name="ref_t2", version="aedbd6fe44051c171fd966c280c5c3036f658831", ) def ref_t2( dataframe: pd.DataFrame, split_mask: int, num_features: int, ) -> pd.DataFrame: ... wb = ImperativeWorkflow(name="core.feature_engineering.workflow.fe_wf") wb.add_workflow_input("sqlite_archive", FlyteFile[typing.TypeVar("sqlite")]) sql_task = SQLite3Task( name="dummy.sqlite.task", query_template="select * from data", inputs=kwtypes(), output_schema_type=FlyteSchema, task_config=SQLite3Config( uri="https://sample/data", compressed=True, ), ) node_sql = wb.add_task(sql_task) node_t1 = wb.add_task(ref_t1, dataframe=node_sql.outputs["results"], imputation_method="mean") node_t2 = wb.add_task( ref_t2, dataframe=node_t1.outputs["o0"], split_mask=24, num_features=15, ) wb.add_workflow_output("output_from_t3", node_t2.outputs["o0"], python_type=pd.DataFrame) wf_spec = get_serializable(OrderedDict(), serialization_settings, wb) assert len(wf_spec.template.nodes) == 3 assert len(wf_spec.template.interface.inputs) == 1 assert wf_spec.template.interface.inputs["sqlite_archive"].type.blob is not None assert len(wf_spec.template.interface.outputs) == 1 assert wf_spec.template.interface.outputs["output_from_t3"].type.structured_dataset_type is not None assert wf_spec.template.interface.outputs["output_from_t3"].type.structured_dataset_type == StructuredDatasetType( format="parquet" )
def test_execute_sqlite3_task(flyteclient, flyte_workflows_register, flyte_remote_env): remote = FlyteRemote(Config.auto(), PROJECT, "development") example_db = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" interactive_sql_task = SQLite3Task( "basic_querying", query_template= "select TrackId, Name from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), output_schema_type=FlyteSchema[kwtypes(TrackId=int, Name=str)], task_config=SQLite3Config( uri=example_db, compressed=True, ), ) registered_sql_task = remote.register(interactive_sql_task) execution = remote.execute(registered_sql_task, inputs={"limit": 10}, wait=True) output = execution.outputs["results"] result = output.open().all() assert result.__class__.__name__ == "DataFrame" assert "TrackId" in result assert "Name" in result
def test_task_serialization(): sql_task = SQLite3Task( "test", query_template= "select TrackId, Name from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), output_schema_type=FlyteSchema[kwtypes(TrackId=int, Name=str)], task_config=SQLite3Config( uri=EXAMPLE_DB, compressed=True, ), ) tt = sql_task.serialize_to_model(sql_task.SERIALIZE_SETTINGS) assert tt.container.args == [ "pyflyte-execute", "--inputs", "{{.input}}", "--output-prefix", "{{.outputPrefix}}", "--raw-output-data-prefix", "{{.rawOutputDataPrefix}}", "--resolver", "flytekit.core.python_customized_container_task.default_task_template_resolver", "--", "{{.taskTemplatePath}}", "flytekit.extras.sqlite3.task.SQLite3TaskExecutor", ] assert tt.custom[ "query_template"] == "select TrackId, Name from tracks limit {{.inputs.limit}}" assert tt.container.image != ""
def test_task_static(): tk = SQLite3Task( "test", query_template="select * from tracks", task_config=SQLite3Config(uri=EXAMPLE_DB, compressed=True,), ) assert tk.output_columns is None df = tk() assert df is not None
def query_wf() -> int: df = SQLite3Task( name="cookbook.sqlite3.sample_inline", query_template= "select TrackId, Name from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), output_schema_type=FlyteSchema[kwtypes(TrackId=int, Name=str)], task_config=SQLite3Config(uri=EXAMPLE_DB, compressed=True), )(limit=100) return print_and_count_columns(df=df)
def test_task_schema(): sql_task = SQLite3Task( "test", query_template="select TrackId, Name from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), output_schema_type=FlyteSchema[kwtypes(TrackId=int, Name=str)], task_config=SQLite3Config(uri=EXAMPLE_DB, compressed=True,), ) assert sql_task.output_columns is not None df = sql_task(limit=1) assert df is not None
def test_workflow(): @task def my_task(df: pandas.DataFrame) -> int: return len(df[df.columns[0]]) sql_task = SQLite3Task( "test", query_template="select * from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), task_config=SQLite3Config(uri=EXAMPLE_DB, compressed=True,), ) @workflow def wf(limit: int) -> int: return my_task(df=sql_task(limit=limit)) assert wf(limit=5) == 5
def test_task_schema(): # sqlite3_start DB_LOCATION = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" sql_task = SQLite3Task( "test", query_template= "select TrackId, Name from tracks limit {{.inputs.limit}}", inputs=kwtypes(limit=int), output_schema_type=FlyteSchema[kwtypes(TrackId=int, Name=str)], task_config=SQLite3Config( uri=DB_LOCATION, compressed=True, ), ) # sqlite3_end assert sql_task.output_columns is not None df = sql_task(limit=1) assert df is not None
except client.exceptions.BucketAlreadyOwnedByYou: logger.info(f"Bucket {bucket_name} has already been created by you.") pass return bucket_name # %% # This is the first task and represents the data source. This can be any task, that fetches data, generates, modifies # data ready for feature ingestion. These can also be arbitrary feature engineering tasks like data imputation, univariate # selection, etc. load_horse_colic_sql = SQLite3Task( name="sqlite3.load_horse_colic", query_template="select * from data", output_schema_type=FlyteSchema, task_config=SQLite3Config( uri=DATABASE_URI, compressed=True, ), metadata=TaskMetadata( cache=True, cache_version="1.0", ), ) # %% # We define two tasks, namely ``store_offline`` and ``load_historical_features`` to store and retrieve the historial # features. # # .. list-table:: Decoding the ``Feast`` Nomenclature # :widths: 25 25 #
import pandas from flytekit import kwtypes, task, workflow from flytekit.extras.sqlite3.task import SQLite3Config, SQLite3Task # https://www.sqlitetutorial.net/sqlite-sample-database/ from flytekit.types.schema import FlyteSchema EXAMPLE_DB = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip" # This task belongs to test_task_static but is intentionally here to help test tracking tk = SQLite3Task( "test", query_template="select * from tracks", task_config=SQLite3Config( uri=EXAMPLE_DB, compressed=True, ), ) def test_task_static(): assert tk.output_columns is None df = tk() assert df is not None def test_task_schema(): # sqlite3_start DB_LOCATION = "https://www.sqlitetutorial.net/wp-content/uploads/2018/03/chinook.zip"
name="great_expectations_task_schema", datasource_name="data", inputs=kwtypes(dataset=FlyteSchema), expectation_suite_name="sqlite.movies", data_connector_name="data_flytetype_data_connector", local_file_path="/tmp/test.parquet", context_root_dir=CONTEXT_ROOT_DIR, ) # %% # Let's fetch the DataFrame from the SQL Database we've with us. To do so, we use the ``SQLite3Task`` available within Flyte. sql_to_df = SQLite3Task( name="greatexpectations.task.sqlite3", query_template="select * from movies", output_schema_type=FlyteSchema, task_config=SQLite3Config(uri=SQLITE_DATASET), ) # %% # Next, we define a task that validates the data and returns the columns in it. @task(limits=Resources(mem="500Mi")) def schema_task(dataset: pd.DataFrame) -> typing.List[str]: schema_task_object(dataset=dataset) return list(dataset.columns) # %% # Finally, we define a workflow to fetch the DataFrame and validate it. @workflow def schema_wf() -> typing.List[str]: df = sql_to_df()