def testApplyToSequenceFuntion(self):
def my_func(row):
return row.name + str(row.id)
self.engine.compile(self.expr.apply(my_func, axis=1, reduce=True).rename('test'))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf, globals(), locals())
udf = locals()[UDF_CLASS_NAME]
self.assertEqual(['name1', 'name2'],
runners.simple_run(udf, [('name', 1, None), ('name', 2, None)]))
def testApplyFunction(self):
def my_func(row):
return row.name, row.id
self.engine.compile(self.expr.apply(my_func, axis=1, names=['name', 'id'], types=['string', 'int']))
udtf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udtf, globals(), locals())
udtf = locals()[UDF_CLASS_NAME]
self.assertEqual([('name1', 1), ('name2', 2)],
runners.simple_run(udtf, [('name1', 1, None), ('name2', 2, None)]))
def testApplyGeneratorFunction(self):
def my_func(row):
for n in row.name.split(','):
yield n
self.engine.compile(self.expr.apply(my_func, axis=1, names='name'))
udtf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udtf, globals(), locals())
udtf = locals()[UDF_CLASS_NAME]
self.assertEqual(['name1', 'name2', 'name3', 'name4'],
runners.simple_run(udtf, [('name1,name2', 1, None), ('name3,name4', 2, None)]))
def testSimpleFunction(self):
def my_func(x):
if x < 0:
return -1
elif x == 0:
return 0
else:
return 1
self.engine.compile(self.expr.id.map(my_func))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf)
udf = locals()[UDF_CLASS_NAME]
self.assertSequenceEqual([-1, 0, 1], runners.simple_run(udf, [(-3, ), (0, ), (5, )]))
def testGlobalVarFunction(self):
global_val = 10
def my_func(x):
if x < global_val:
return -1
elif x == global_val:
return 0
else:
return 1
self.engine.compile(self.expr.id.map(my_func))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf, globals(), locals())
udf = locals()[UDF_CLASS_NAME]
self.assertSequenceEqual([-1, 0, 1], runners.simple_run(udf, [(-9, ), (10, ), (15, )]))
def testNestFunction(self):
def my_func(x):
def inner(y):
if y < 0:
return -2
elif y == 0:
return 0
else:
return 2
return inner(x)
self.engine.compile(self.expr.id.map(my_func))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf, globals(), locals())
udf = locals()[UDF_CLASS_NAME]
self.assertSequenceEqual([-2, 0, 2], runners.simple_run(udf, [(-3, ), (0, ), (5, )]))
def testRefFuncFunction(self):
global_val = 10
def my_func1(x):
if x < global_val:
return -1
elif x == global_val:
return 0
else:
return 1
def my_func(y):
return my_func1(y)
self.engine.compile(self.expr.id.map(my_func))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf, globals(), locals())
udf = locals()[UDF_CLASS_NAME]
self.assertSequenceEqual([-1, 0, 1],
runners.simple_run(udf, [(-9, ), (10, ),
(15, )]))
def testAggNumpyFunction(self):
import numpy as np
class Agg(object):
def buffer(self):
return [np.int32(1)]
def __call__(self, buffer, val):
buffer[0] *= val
def merge(self, buffer, pbuffer):
buffer[0] *= pbuffer[0]
def getvalue(self, buffer):
return buffer[0]
self.engine.compile(self.expr.id.agg(Agg))
udaf = list(self.engine._ctx._func_to_udfs.values())[0]
udaf = get_function(udaf, UDF_CLASS_NAME)
result = runners.simple_run(udaf, [(3,), (6,), (5,)])
self.assertEqual([90], result)
self.assertNotIsInstance(result, np.generic)
def testSimpleLambda(self):
self.engine.compile(self.expr.id.map(lambda x: x + 1))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
six.exec_(udf)
udf = locals()[UDF_CLASS_NAME]
self.assertSequenceEqual([4, ], runners.simple_run(udf, [(3, ), ]))
def testSimpleLambda(self):
self.engine.compile(self.expr.id.map(lambda x: x + 1))
udf = list(self.engine._ctx._func_to_udfs.values())[0]
udf = get_function(udf, UDF_CLASS_NAME)
self.assertSequenceEqual([4, ], runners.simple_run(udf, [(3, ), ]))
def test_udtf(self):
self.assertEqual(['a', 'b', 'ok'], runners.simple_run(Explode, [('a|b',),]))
def test_udaf(self):
self.assertEqual([2], runners.simple_run(Avg, [(1,),(2,),(3,)]))
def test_udf(self):
self.assertEqual([2,3], runners.simple_run(Plus, [(1,1), (2,1)]))
self.assertEqual([None], runners.simple_run(Plus, [(None,1) ]))
def test_udtf(self):
self.assertEqual(["a", "b", "ok"], runners.simple_run(Explode, [("a|b",)]))
def test_udaf(self):
self.assertEqual([2], runners.simple_run(Avg, [(1,), (2,), (3,)]))
def test_udf(self):
self.assertEqual([2, 3], runners.simple_run(Plus, [(1, 1), (2, 1)]))
self.assertEqual([None], runners.simple_run(Plus, [(None, 1)]))
