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_base64_encode_tostring():
df = pd.DataFrame()
df["A"] = [_encode_base64("hello"), _encode_base64("There ")]
h = "hello"
w = Wrap(df)
w = w.let(h=h).extend("B = A == base64_encode_tostring(h)")
assert [True, False] == list(w.df["B"])
def test_execute_let2_pipe():
df = create_df()
w = Wrap(df)
wnew = w.let(
a='G1').execute("self | where G == a | where A >= 1 | project C")
assert ["C"] == list(wnew.df.columns)
assert ["foo2", "foo4"] == list(wnew.df["C"])
def test_let_new_method_overrides_old_one():
df = pd.DataFrame()
df["A"] = [1, 2]
w = Wrap(df)
w = w.let(b=1)
w = w.let(b=2)
wnew = w.extend("B = A + b")
assert ["A", "B"] == list(wnew.df.columns)
assert [3, 4] == list(wnew.df["B"])
def test_let_no_params():
df = pd.DataFrame()
df["A"] = [1, 1]
w = Wrap(df)
def x():
return 3
w = w.let(x=x)
wnew = w.extend("B = A + x()")
assert [4, 4] == list(wnew.df["B"])
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_let_method(self):
df = create_df()
w = Wrap(df)
wnew = w.let(xx=lambda x: x + 1).extend("Z = B + xx(4)")
self.assertListEqual([5, 6, 7, 8, 9], list(wnew.df["Z"]))
def test_let(self):
df = create_df()
w = Wrap(df)
wnew = w.let(xx=2).let(yy=3).extend("Z = B + xx + yy")
self.assertListEqual([5, 6, 7, 8, 9], list(wnew.df["Z"]))
