def test_union_execute_multiple_tables():
df = pd.DataFrame()
df["A"] = [1, 2]
df["B"] = [10, 20]
df2 = pd.DataFrame()
df2["A"] = [3, 4]
df2["C"] = [30, 40]
df3 = pd.DataFrame()
df3["A"] = [5, 6]
df3["D"] = [50, 60]
w = Wrap(df)
w = w.let(df2=df2, df3=df3)
wnew = w.execute("""
self
| union kind=outer df2, df3
""")
assert ["A", "B", "C", "D"] == list(wnew.df.columns)
assert [1, 2, 3, 4, 5, 6] == list(wnew.df["A"])
assert [10, 20, -1, -1, -1, -1] == replace_nan(wnew.df["B"], -1)
assert [-1, -1, 30, 40, -1, -1] == replace_nan(wnew.df["C"], -1)
assert [-1, -1, -1, -1, 50, 60] == replace_nan(wnew.df["D"], -1)
def test_execute_let_semicolon():
df = create_df()
w = Wrap(df)
wnew = w.execute("let a = 'G2';")
assert "G2" == wnew._get_var_map()["a"]
def test_execute():
df = create_df()
w = Wrap(df)
wnew = w.execute("self | where G == 'G1' | where A >= 1 | project C")
assert ["C"] == list(wnew.df.columns)
assert ["foo2", "foo4"] == list(wnew.df["C"])
def test_execute_wrap_object():
df = create_df()
w = Wrap(df)
w2 = Wrap(pd.DataFrame()).let(w=w)
wnew = w2.execute("w | where G == 'G1' | where A >= 1 | project C")
assert ["C"] == list(wnew.df.columns)
assert ["foo2", "foo4"] == list(wnew.df["C"])
def test_execute_join():
df = create_df()
w = Wrap(df)
wnew = w.execute(
"self | where G == 'G1' | project G | as T1; self | join kind=inner (T1) on G | project A, G"
)
assert ["A", "G"] == list(wnew.df.columns)
assert ["G1"] * 9 == list(wnew.df["G"])
def test_execute_let_dataframe():
df = create_df()
df2 = pd.DataFrame()
w = Wrap(df2).let(D=df)
wnew = w.execute(
"let a = 'G1'; D | where G == a | where A >= 1 | project C")
assert ["C"] == list(wnew.df.columns)
assert ["foo2", "foo4"] == list(wnew.df["C"])
def test_order_execute(self):
df = create_df()
df["U"] = [9, 1, 7, 1, 2]
expected = [1, 3, 4, 2, 0]
w = Wrap(df)
wnew = w.execute("self | order by U + 1 asc, B asc")
self.assertListEqual([1, 1, 2, 7, 9], list(wnew.df["U"]))
self.assertListEqual(expected, list(wnew.df["B"]))
self.assertListEqual(list(range(5)), list(w.df["B"]))
assert 6 == len(wnew.df.columns)
def test_execute_comment():
df = create_df()
w = Wrap(df)
wnew = w.execute("""
self #c0
| where G == 'G1' # c1
# comment
| where A >= 1 #c3
| project C # another comment
""")
assert ["C"] == list(wnew.df.columns)
assert ["foo2", "foo4"] == list(wnew.df["C"])
def test_execute_join_right_on_left_on():
df = create_df()
w = Wrap(df)
wnew = w.execute("""
self
| join (
self
| where G == 'G1'
| project GG=G )
on $right.GG == $left.G
| project A, G
""")
assert ["A", "G"] == list(wnew.df.columns)
assert ["G1"] * 9 == list(wnew.df["G"])
def test_execute_join_nokind():
df = create_df()
w = Wrap(df)
wnew = w.execute("""
self
| join (
self
| where G == 'G1'
| project G )
on G
| project A, G
""")
assert ["A", "G"] == list(wnew.df.columns)
assert ["G1"] * 9 == list(wnew.df["G"])
def test_union_execute():
df = pd.DataFrame()
df["A"] = [1, 2]
df["B"] = [10, 20]
df2 = pd.DataFrame()
df2["A"] = [3, 4]
df2["C"] = [30, 40]
w = Wrap(df)
w = w.let(df2=df2)
wnew = w.execute("""
self
| union kind=outer df2
""")
assert ["A", "B", "C"] == list(wnew.df.columns)
assert [1, 2, 3, 4] == list(wnew.df["A"])
assert [10, 20, -1, -1] == replace_nan(wnew.df["B"], -1)
assert [-1, -1, 30, 40] == replace_nan(wnew.df["C"], -1)
def test_execute_join_multiple_join_conditions():
df = pd.DataFrame()
df["A"] = [1, 1]
df["B"] = [1, 2]
df["C"] = [10, 20]
df["D"] = [100, 200]
df2 = pd.DataFrame()
df2["A"] = [1, 1]
df2["B2"] = [1, 0]
df2["C2"] = [10, 0]
w = Wrap(df)
w = w.let(df2=df2)
wnew = w.execute("""
self
| join ( df2 )
on A, $right.B2 == $left.B, $left.C == $right.C2
""")
assert ["A", "B", "C", "D", "B2", "C2"] == list(wnew.df.columns)
assert [100] == list(wnew.df["D"])
def test_limit_expression(self):
df = create_df()
w = Wrap(df)
wnew = w.execute("self | limit 1+1")
self.assertListEqual([0, 1], list(wnew.df["B"]))
self.assertEqual(5, len(w.df["B"]))
