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, ), ]))
Beispiel #11
0
 def test_udtf(self):
     self.assertEqual(['a', 'b', 'ok'], runners.simple_run(Explode, [('a|b',),]))
Beispiel #12
0
 def test_udaf(self):
     self.assertEqual([2], runners.simple_run(Avg, [(1,),(2,),(3,)]))
Beispiel #13
0
 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)]))