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)]))