def test_yield(tmpdir): df = pd.DataFrame([[0, 0]], columns=["a", "b"]) # schema: * def t(df: pd.DataFrame) -> pd.DataFrame: return df.assign(b=df.b + 1) dag = FugueWorkflow() dag.df(df).transform(t).yield_dataframe_as("x") result = dag.run()["x"] assert [[0, 1]] == result.as_array() dag1 = FugueWorkflow() dag1.df(df).transform(t).yield_file_as("x") dag1.run("", {FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)}) dag2 = FugueWorkflow() dag2.df(dag1.yields["x"]).transform(t).yield_dataframe_as("y") result = dag2.run("", {FUGUE_CONF_WORKFLOW_CHECKPOINT_PATH: str(tmpdir)})["y"] assert [[0, 2]] == result.as_array() dag3 = FugueWorkflow() dag3.df(dag2.yields["y"]).transform(t).yield_dataframe_as("z") result = dag3.run()["z"] assert [[0, 3]] == result.as_array()
def test_worflow_dataframes(): dag1 = FugueWorkflow() df1 = dag1.df([[0]], "a:int") df2 = dag1.df([[0]], "b:int") dag2 = FugueWorkflow() df3 = dag2.df([[0]], "a:int") dfs1 = WorkflowDataFrames(a=df1, b=df2) assert dfs1["a"] is df1 assert dfs1["b"] is df2 dfs2 = WorkflowDataFrames(dfs1, aa=df1, bb=df2) assert 4 == len(dfs2) with raises(ValueError): WorkflowDataFrames(a=df1, b=df3) with raises(ValueError): WorkflowDataFrames(a=df1, b=ArrayDataFrame([[0]], "a:int")) dag = FugueWorkflow() df = dag.df([[0], [1]], "a:int") assert df.partition_spec.empty df2 = df.partition(by=["a"]) assert df.partition_spec.empty assert df2.partition_spec == PartitionSpec(by=["a"])
def test_process(): # basic features, nest dag = FugueWorkflow() a1 = dag.create(mock_create1, params=dict(n=1)) a2 = dag.create(mock_create1, params=dict(n=2)) dag.process(a1, a2, using=mock_processor1, params=dict(n=3)) dag.process(a2, a1, using=mock_processor2, schema="b:int", params=dict(n=4)) dag.process( dag.create(mock_create1, params=dict(n=5)), dag.create(mock_create1, params=dict(n=6)), using=mock_processor1, params=dict(n=7)) assert_eq(""" a=create using mock_create1 params n:1 b=create using mock_create1 params n:2 process a,b using mock_processor1(n=3) process b,a using mock_processor2(n=4) schema b:int process # nested (create using mock_create1(n=5)), (create using mock_create1(n=6)) using mock_processor1(n=7) """, dag) # anonymous, nested anonymous dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)).process(mock_processor3) b = a.partition(by=["a"]).process(mock_processor3) c = a.process(mock_processor3) dag.process(b, c, using=mock_processor1) assert_eq(""" create using mock_create1 params n:1 process using mock_processor3 process # nested (process prepartition by a using mock_processor3), (process using mock_processor3) using mock_processor1 """, dag) # no last dataframe with raises(FugueSQLError): assert_eq(""" process using mock_processor3 """, None) # dict like dataframes dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)) b = dag.create(mock_create1, params=dict(n=2)) dag.process(dict(df1=a, df2=b), using=mock_processor1) assert_eq(""" process df1=(create using mock_create1(n=1)), df2:(create using mock_create1(n=2)) using mock_processor1 """, dag)
def test_workflow(): builder = FugueWorkflow() a = builder.create_data([[0], [0], [1]], "a:int") raises(InvalidOperationError, lambda: a._task.copy()) raises(InvalidOperationError, lambda: copy.copy(a._task)) raises(InvalidOperationError, lambda: copy.deepcopy(a._task)) a.show() a.show() raises(FugueWorkflowCompileError, lambda: builder.df(123)) b = a.transform(mock_tf1, "*,b:int", pre_partition=dict(by=["a"])) b.show() builder.create_data([[0], [1]], "b:int").show() c = ArrayDataFrame([[100]], "a:int") builder.show(a, b, c) b = a.partition(by=["a"]).transform(mock_tf2).persist().broadcast() b.show() builder.run() df_eq(a.result, [[0], [0], [1]], "a:int") raises(TypeError, lambda: builder.run("abc")) builder.run(FugueWorkflowContext()) df_eq(a.result, [[0], [0], [1]], "a:int") builder.run("NativeExecutionEngine") df_eq(b.result, [[0, 2], [0, 2], [1, 1]], "a:int,b:int") df_eq(b.compute(), [[0, 2], [0, 2], [1, 1]], "a:int,b:int") df_eq(b.compute(NativeExecutionEngine), [[0, 2], [0, 2], [1, 1]], "a:int,b:int")
def test_workflows(self): a = FugueWorkflow().df([[0]], "a:int") df_eq(a.compute(self.engine), [[0]], "a:int") a = _FugueInteractiveWorkflow(self.engine).df([[0]], "a:int").persist() df_eq(a.result, [[0]], "a:int")
def test_create(): dag = FugueWorkflow() dag.create(mock_create1, params=dict(n=1)) dag.create(mock_create2, schema="a:int", params=dict(n=1)) assert_eq(""" a=create using mock_create1 params n:1 b=create using mock_create2(n=1) schema a:int """, dag)
def test_select(): dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)) b = dag.create(mock_create1, params=dict(n=2)) dag.select("select * from a.b") dag.select("select * from a.b TABLESAMPLE (5 PERCENT) AS x") dag.select("select * from a.b AS x") dag.select("select * from", a, "AS a") # fugue sql adds 'AS a' dag.select("select * from", a, "TABLESAMPLE (5 PERCENT) AS a") x = dag.select("select * from", a, "TABLESAMPLE (5 PERCENT) AS x") y = dag.select("select * FROM", x) z = dag.select("select * FROM", y, "where t = 100") dag.select("select a.* from", a, "AS a join", b, "AS b on a.a == b.a") dag.select("select * from (select * from a.b)") dag.select("select * from", dag.create(mock_create1), "TABLESAMPLE (5 PERCENT)") dag.select("select * from", dag.create(mock_create1), "AS b") dag.select("select * from (select * from", dag.create(mock_create1), ")") dag.select("select * from", a, "AS a").persist().broadcast().show() dag.select("select * from", a, "AS a").persist("a.b.c").broadcast().show() assert_eq(""" a=create using mock_create1(n=1) b=create using mock_create1(n=2) # assignment and table not found x=select * from a.b # sample and alias when table not found select * from a.b TABLESAMPLE (5 PERCENT) AS x select * from a.b AS x # when table is found select * from a select * from a TABLESAMPLE(5 PERCENT) select * from a TABLESAMPLE(5 PERCENT) AS x # no from select * select * where t=100 # multiple dependencies select a.* from a join b on a.a==b.a # nested query select * from (select * from a.b) # nested fugue extensions select * from (create using mock_create1) TABLESAMPLE(5 PERCENT) select * from (create using mock_create1) AS b select * from (select * from (create using mock_create1)) # persist & checkpoint & broadcast select * from a persist broadcast print select * from a persist a.b.c broadcast print """, dag)
def test_zip(): dag = FugueWorkflow() a1 = dag.create(mock_create1, params=dict(n=1)) a2 = dag.create(mock_create1, params=dict(n=2)) a1.zip(a2) assert_eq(""" a=create using mock_create1 params n:1 zip a,(create using mock_create1 params n:2) """, dag) dag = FugueWorkflow() a1 = dag.create(mock_create1, params=dict(n=1)) a2 = dag.create(mock_create1, params=dict(n=2)) a1.zip(a2, how="left_outer", partition=dict(by=["a"], presort="b DESC")) assert_eq(""" a=create using mock_create1 params n:1 zip a,(create using mock_create1 params n:2) left outer by a presort b desc """, dag)
def test_transform(): w = (FugueWorkflow().df([[0], [1]], "a:int") .transform(mock_transformer, schema=Schema("a:int"), params=dict(n=2)) ) assert_eq(""" create [[0],[1]] schema a:int transform using mock_transformer(n=2) schema a:int """, w.workflow) w = (FugueWorkflow().df([[0], [1]], "a:int") .partition(by=["a"], presort="b DESC", num="ROWCOUNT/2") .transform(mock_transformer, schema="*", params=dict(n=2)) ) assert_eq(""" create [[0],[1]] schema a:int transform prepartition ROWCOUNT / 2 by a presort b desc using mock_transformer(n=2) schema * """, w.workflow)
def test_print(): dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)) a.show() b = dag.create(mock_create1, params=dict(n=2)) dag.show(a, b, rows=5, show_count=True, title="\"b B") assert_eq(""" a=create using mock_create1(n=1) print print a, (create using mock_create1(n=2)) rows 5 rowcount title "\\"b B" """, dag)
def test_output(): dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)) a.partition(num=4).output(mock_output) b = dag.create(mock_create1, params=dict(n=2)) dag.output(a, b, using=mock_output, params=dict(n=3)) assert_eq(""" a=create using mock_create1(n=1) output prepartition 4 using mock_output output a, (create using mock_create1(n=2)) using mock_output(n=3) """, dag)
def test_dataset(tmpdir): factory = TuneObjectFactory() factory.set_temp_path(str(tmpdir)) df = pd.DataFrame([[0]], columns=["a"]) dag = FugueWorkflow() data = factory.make_dataset(dag, Space(a=1, b=1), df=df) assert isinstance(data, TuneDataset) assert factory.make_dataset(dag, data) is data # TODO: the unit test is not complete, but it's covered by other functions with raises(TuneCompileError): factory.make_dataset(dag, 1)
def test_load(): dag = FugueWorkflow() dag.load("xx") dag.load("xx", fmt="csv") dag.load("xx", columns="a:int,b:str") dag.load("xx", columns=["a", "b"], header=True) assert_eq(""" load "xx" load csv "xx" load "xx" columns a:int, b:str load "xx"(header=True) columns a, b """, dag)
def test_cotransform(): dag = FugueWorkflow() a1 = dag.create(mock_create1, params=dict(n=1)) a2 = dag.create(mock_create1, params=dict(n=2)) z = dag.zip(a1, a2) t = z.partition(num=3).transform(mock_cotransformer1, params=dict(n=3)) assert_eq(""" zip (create using mock_create1 params n:1), (create using mock_create1 params n:2) transform prepartition 3 using mock_cotransformer1(n=3) """, dag)
def test_save(): dag = FugueWorkflow() a = dag.create(mock_create1, params=dict(n=1)) a.save("xx", fmt="parquet", mode="overwrite") a.save("xx", mode="append") a.save("xx", mode="error") a.save("xx.csv", fmt="csv", mode="error", single=True, header=True) a.partition(by=["x"]).save("xx", mode="overwrite") b = dag.create(mock_create1, params=dict(n=2)).save("xx", mode="overwrite") assert_eq(""" a=create using mock_create1(n=1) save overwrite parquet "xx" save a append "xx" save a to "xx" save to single csv "xx.csv"(header=True) save prepartition by x overwrite "xx" save (create using mock_create1(n=2)) overwrite "xx" """, dag)
def test_iterative_study(tmpdir): def assert_metric(df: Iterable[Dict[str, Any]], metric: float) -> None: for row in df: assert row[TUNE_REPORT_METRIC] < metric study = IterativeStudy(F(), str(tmpdir)) space = sum( Space(a=a, b=b) for a, b in [(1.1, 0.2), (0.8, -0.2), (1.2, -0.1), (0.7, 0.3), (1.0, 1.5)]) dag = FugueWorkflow() dataset = TuneDatasetBuilder(space, str(tmpdir)).build(dag) result = study.optimize( dataset, J([1, 2, 3, 4]), ) result.result(1).show() result.result(1).output(assert_metric, params=dict(metric=-2.8)) dag.run()
def test_persist_checkpoint_broadcast(): dag = FugueWorkflow() dag.create(mock_create1).persist() dag.create(mock_create1).persist("a.b") dag.create(mock_create1).broadcast() dag.create(mock_create1).persist("a.b").broadcast() dag.create(mock_create1).checkpoint() dag.create(mock_create1).checkpoint() dag.create(mock_create1).checkpoint("xy z") dag.create(mock_create1).checkpoint("xy z").broadcast() assert_eq(""" create using mock_create1 persist a=create using mock_create1 persist a.b create using mock_create1 broadcast a=create using mock_create1 persist a.b broadcast create using mock_create1 checkpoint a?? create using mock_create1 a=create using mock_create1 checkpoint "xy z" a??create using mock_create1 checkpoint "xy z" broadcast """, dag)
def test_default_init(self): a = FugueWorkflow().df([[0]], "a:int") df_eq(a.compute(DaskExecutionEngine), [[0]], "a:int")
def test_create_data(): w = FugueWorkflow().df([[0], [1]], "a:int") assert_eq(""" a=create [[0],[1]] schema a:int """, w.workflow)
def dag(self) -> FugueWorkflow: return FugueWorkflow(self.engine)
def test_default_session(self): a = FugueWorkflow().df([[0]], "a:int") df_eq(a.compute(SparkExecutionEngine), [[0]], "a:int")
def test_df_init(self): sdf = self.spark_session.createDataFrame([[1.1]], "a:double") a = FugueWorkflow().df(sdf) df_eq(a.compute(SparkExecutionEngine), [[1.1]], "a:double")
def assert_eq(expr, expected: FugueWorkflow): sql = FugueSQL(expr, "fugueLanguage", ignore_case=True, simple_assign=True) wf = FugueWorkflow() v = _Extensions(sql, FugueSQLHooks(), wf) obj = v.visit(sql.tree) assert expected.spec_uuid() == v.workflow.spec_uuid()
def test_session_as_engine(self): dag = FugueWorkflow() a = dag.df([[p, 0] for p in range(100)], "a:int,b:int") a.partition(algo="even", by=["a"]).transform(AssertMaxNTransform).persist() dag.run(self.spark_session)