def test_next_reset_combined():
def udf_wrapper():
def run(ctx):
for i in range(2):
ctx.emit(ctx.t)
if not ctx.next():
break
ctx.next(reset=True)
for i in range(2):
ctx.emit(ctx.t+1)
if not ctx.next():
break
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")]
)
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1,), (5,), (6,)])], exa)
assert result == [Group([(1,), (5,),(2,), (6,)])]
def test_udf_wrapper_with_no_empty_line_after_function_name():
def udf_wrapper():
def run(ctx):
pass
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, )])], exa)
assert result == [Group([])]
def test_exception_udf_wrapper_with_parameter():
def udf_wrapper(param):
def run(ctx):
pass
executor = UDFMockExecutor()
with pytest.raises(Exception):
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, )])], exa)
def test_emit_single_column_none():
def udf_wrapper():
def run(ctx):
ctx.emit(None)
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, ), (5, ), (6, )])], exa)
assert result == [Group([(None, )])]
def test_get_dataframe_start_col_negative():
def udf_wrapper():
def run(ctx):
df = ctx.get_dataframe(num_rows=10, start_col=-1)
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
with pytest.raises(RuntimeError) as excinfo:
result = executor.run(
[Group([(1, ), (2, ), (3, ), (4, ), (5, ), (6, )])], exa)
def test_get_dataframe_all():
def udf_wrapper():
def run(ctx):
df = ctx.get_dataframe(num_rows='all')
ctx.emit(df)
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, ), (5, ), (6, )])], exa)
assert result == [Group([(1, ), (5, ), (6, )])]
def test_emit_tuple_exception():
def udf_wrapper():
def run(ctx):
while True:
ctx.emit((1, ))
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
with pytest.raises(TypeError):
result = executor.run(
[Group([(1, ), (2, ), (3, ), (4, ), (5, ), (6, )])], exa)
def test_exa_meta_in_run():
def udf_wrapper():
def run(ctx):
ctx.emit(exa.meta.script_code)
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", str, "VARCHAR(2000)")]
)
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1,)])], exa)
assert result == [Group([(exa.meta.script_code,)])]
def test_emit_not_allowed():
def udf_wrapper():
def run(ctx):
ctx.emit(ctx.t)
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SCALAR",
input_columns=[Column("t", int, "INTEGER")],
output_type="RETURNS",
output_columns=[Column("t", int, "INTEGER")]
)
exa = MockExaEnvironment(meta)
with pytest.raises(RuntimeError):
result = executor.run([Group([(1,), (5,), (6,)])], exa)
def test_simple_return():
def udf_wrapper():
def run(ctx):
return ctx.t+1
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SCALAR",
input_columns=[Column("t", int, "INTEGER")],
output_type="RETURNS",
output_columns=[Column("t", int, "INTEGER")]
)
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1,), (5,), (6,)])], exa)
assert result == [Group([(2,), (6,), (7,)])]
def test_next_and_emit():
def udf_wrapper():
def run(ctx):
while True:
ctx.emit(ctx.t)
if not ctx.next():
return
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, ), (5, ), (6, )])], exa)
assert result == [Group([(1, ), (5, ), (6, )])]
def test_return_multi_column_none():
def udf_wrapper():
def run(ctx):
return None,None
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SCALAR",
input_columns=[Column("t", int, "INTEGER")],
output_type="RETURNS",
output_columns=[Column("t1", int, "INTEGER"),
Column("t2", int, "INTEGER")]
)
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1,), (5,), (6,)])], exa)
assert result == [Group([(None,None),(None,None),(None,None)])]
def test_get_connection_in_init():
def udf_wrapper():
con = exa.get_connection("TEST_CON")
def run(ctx):
ctx.emit(con.address)
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", str, "VARCHAR(2000)")]
)
exa = MockExaEnvironment(meta, connections={"TEST_CON": Connection(address="https://test.de")})
result = executor.run([Group([(1,)])], exa)
assert result == [Group([("https://test.de",)])]
def test_udf_wrapper_with_docstring_after_empty_lines():
def udf_wrapper():
"""
wrapper with docstring should raise Exception,
because their is no easy way to remove docstrings
to get only the source witin the function
"""
def run(ctx):
pass
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, )])], exa)
assert result == [Group([])]
def test_multi_column_type():
def udf_wrapper():
def run(ctx):
return ctx.t1+1, ctx.t2+1.1, ctx.t3+"1"
executor = UDFMockExecutor()
meta = MockMetaData(
script_code_wrapper_function=udf_wrapper,
input_type="SCALAR",
input_columns=[Column("t1", int, "INTEGER"),
Column("t2", float, "FLOAT"),
Column("t3", str, "VARCHAR(20000)")],
output_type="RETURNS",
output_columns=[Column("t1", int, "INTEGER"),
Column("t2", float, "FLOAT"),
Column("t3", str, "VARCHAR(20000)")]
)
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1,1.0,"1"), (5,5.0,"5"), (6,6.0,"6")])], exa)
assert result == [Group([(2,2.1,"11"), (6,6.1,"51"), (7,7.1,"61")])]
def test_get_dataframe_iter_next():
def udf_wrapper():
def run(ctx):
while True:
df = ctx.get_dataframe(num_rows=2)
if df is None:
return
else:
ctx.emit(df)
ctx.next()
executor = UDFMockExecutor()
meta = MockMetaData(script_code_wrapper_function=udf_wrapper,
input_type="SET",
input_columns=[Column("t", int, "INTEGER")],
output_type="EMITS",
output_columns=[Column("t", int, "INTEGER")])
exa = MockExaEnvironment(meta)
result = executor.run([Group([(1, ), (2, ), (3, ), (4, ), (5, ), (6, )])],
exa)
assert result == [Group([(1, ), (2, ), (4, ), (5, )])]