def test_lambda_in_filter():
df = DataFrame()
df1 = df[df.apply(lambda e: e == 1)]
assert isinstance(df1.child_expr, ast_DataFrame)
assert df1.filter is not None
assert isinstance(df1.filter, Column)
assert isinstance(df1.filter.child_expr, ast.Call)
def test_lambda_argument():
df = DataFrame()
df1 = df.apply(lambda e: e)
assert df1.child_expr is not None
assert isinstance(df1.child_expr, ast.Call)
assert len(df1.child_expr.args) == 1
arg1 = df1.child_expr.args[0]
assert isinstance(arg1, ast_Callable)
def test_callable_wrong_number_args():
d = DataFrame()
d1 = d.apply(lambda b: b)
expr, ctx = render(d1)
assert isinstance(expr, ast.Call)
arg1 = expr.args[0] # type: ast.AST
assert isinstance(arg1, ast_Callable)
with pytest.raises(Exception):
render_callable(arg1, ctx, d, d)
def test_callable_returns_const():
d = DataFrame()
d1 = d.apply(lambda b: 20)
expr, ctx = render(d1)
assert isinstance(expr, ast.Call)
arg1 = expr.args[0] # type: ast.AST
assert isinstance(arg1, ast_Callable)
expr1, new_ctx = render_callable(arg1, ctx, d)
assert isinstance(expr1, ast.Num)
assert expr1.n == 20
def test_callable_simple_call():
d = DataFrame()
d1 = d.apply(lambda b: b)
expr, ctx = render(d1)
assert isinstance(expr, ast.Call)
arg1 = expr.args[0] # type: ast.AST
assert isinstance(arg1, ast_Callable)
expr1, new_ctx = render_callable(arg1, ctx, d)
assert isinstance(expr1, ast_DataFrame)
assert expr1.dataframe is d
def test_callable_function():
def test_func(b):
return b
d = DataFrame()
d1 = d.apply(test_func)
expr, ctx = render(d1)
assert isinstance(expr, ast.Call)
arg1 = expr.args[0] # type: ast.AST
assert isinstance(arg1, ast_Callable)
expr1, new_ctx = render_callable(arg1, ctx, d)
assert isinstance(expr1, ast_DataFrame)
assert expr1.dataframe is d
def test_callable_context_no_update():
d = DataFrame()
d1 = d.apply(lambda b: b.jets.pt)
expr, ctx = render(d1)
assert isinstance(expr, ast.Call)
arg1 = expr.args[0] # type: ast.AST
assert isinstance(arg1, ast_Callable)
seen_ds = len(ctx._seen_datasources)
resolved = len(ctx._resolved)
expr1, new_ctx = render_callable(arg1, ctx, d)
assert len(new_ctx._seen_datasources) != len(ctx._seen_datasources) \
or len(new_ctx._resolved) != len(ctx._resolved)
assert seen_ds == len(ctx._seen_datasources) \
or resolved == len(ctx._resolved)
