| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ################################################################################ |
| import datetime |
| import os |
| import unittest |
| import uuid |
| |
| import pytz |
| |
| from pyflink.common import Row |
| from pyflink.table import DataTypes, expressions as expr |
| from pyflink.table.expressions import call |
| from pyflink.table.udf import ScalarFunction, udf, FunctionContext |
| from pyflink.testing import source_sink_utils |
| from pyflink.testing.test_case_utils import PyFlinkStreamTableTestCase, \ |
| PyFlinkBatchTableTestCase |
| |
| |
| def generate_random_table_name(): |
| return "Table{0}".format(str(uuid.uuid1()).replace("-", "_")) |
| |
| |
| class UserDefinedFunctionTests(object): |
| |
| def test_scalar_function(self): |
| # test metric disabled. |
| self.t_env.get_config().set('python.metric.enabled', 'false') |
| self.t_env.get_config().set('pipeline.global-job-parameters', 'subtract_value:2') |
| # test lambda function |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) |
| |
| # test Python ScalarFunction |
| subtract_one = udf(SubtractOne(), result_type=DataTypes.BIGINT()) |
| |
| subtract_two = udf(SubtractWithParameters(), result_type=DataTypes.BIGINT()) |
| |
| # test callable function |
| add_one_callable = udf(CallablePlus(), result_type=DataTypes.BIGINT()) |
| |
| def partial_func(col, param): |
| return col + param |
| |
| # test partial function |
| import functools |
| add_one_partial = udf(functools.partial(partial_func, param=1), |
| result_type=DataTypes.BIGINT()) |
| |
| # check memory limit is set |
| @udf(result_type=DataTypes.BIGINT()) |
| def check_memory_limit(exec_mode): |
| if exec_mode == "process": |
| assert os.environ['_PYTHON_WORKER_MEMORY_LIMIT'] is not None |
| return 1 |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b BIGINT, c BIGINT, d BIGINT, e BIGINT, f BIGINT, |
| g BIGINT, h BIGINT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| execution_mode = self.t_env.get_config().get("python.execution-mode", "process") |
| |
| t = self.t_env.from_elements([(1, 2, 3), (2, 5, 6), (3, 1, 9)], ['a', 'b', 'c']) |
| t.where(add_one(t.b) <= 3).select( |
| add_one(t.a), |
| subtract_one(t.b), |
| subtract_two(t.b), |
| add(t.a, t.c), |
| add_one_callable(t.a), |
| add_one_partial(t.a), |
| check_memory_limit(execution_mode), |
| t.a).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[2, 1, 0, 4, 2, 2, 1, 1]", "+I[4, 0, -1, 12, 4, 4, 1, 3]"]) |
| |
| def test_chaining_scalar_function(self): |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) |
| subtract_one = udf(SubtractOne(), result_type=DataTypes.BIGINT()) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b BIGINT, c INT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements([(1, 2, 1), (2, 5, 2), (3, 1, 3)], ['a', 'b', 'c']) |
| t.select(add(add_one(t.a), subtract_one(t.b)), t.c, expr.lit(1)) \ |
| .execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[3, 1, 1]", "+I[7, 2, 1]", "+I[4, 3, 1]"]) |
| |
| def test_udf_in_join_condition(self): |
| t1 = self.t_env.from_elements([(2, "Hi")], ['a', 'b']) |
| t2 = self.t_env.from_elements([(2, "Flink")], ['c', 'd']) |
| |
| f = udf(lambda i: i, result_type=DataTypes.BIGINT()) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b STRING, c BIGINT, d StRING) |
| WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| t1.join(t2).where(f(t1.a) == t2.c).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[2, Hi, 2, Flink]"]) |
| |
| def test_udf_in_join_condition_2(self): |
| t1 = self.t_env.from_elements([(1, "Hi"), (2, "Hi")], ['a', 'b']) |
| t2 = self.t_env.from_elements([(2, "Flink")], ['c', 'd']) |
| |
| f = udf(lambda i: i, result_type=DataTypes.BIGINT()) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}( |
| a BIGINT, |
| b STRING, |
| c BIGINT, |
| d STRING |
| ) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t1.join(t2).where(f(t1.a) == f(t2.c)).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[2, Hi, 2, Flink]"]) |
| |
| def test_udf_with_constant_params(self): |
| def udf_with_constant_params(p, null_param, tinyint_param, smallint_param, int_param, |
| bigint_param, decimal_param, float_param, double_param, |
| boolean_param, str_param, |
| date_param, time_param, timestamp_param): |
| from decimal import Decimal |
| import datetime |
| |
| assert null_param is None, 'null_param is wrong value %s' % null_param |
| |
| assert isinstance(tinyint_param, int), 'tinyint_param of wrong type %s !' \ |
| % type(tinyint_param) |
| p += tinyint_param |
| assert isinstance(smallint_param, int), 'smallint_param of wrong type %s !' \ |
| % type(smallint_param) |
| p += smallint_param |
| assert isinstance(int_param, int), 'int_param of wrong type %s !' \ |
| % type(int_param) |
| p += int_param |
| assert isinstance(bigint_param, int), 'bigint_param of wrong type %s !' \ |
| % type(bigint_param) |
| p += bigint_param |
| assert decimal_param == Decimal('1.05'), \ |
| 'decimal_param is wrong value %s ' % decimal_param |
| |
| p += int(decimal_param) |
| |
| assert isinstance(float_param, float) and float_equal(float_param, 1.23, 1e-06), \ |
| 'float_param is wrong value %s ' % float_param |
| |
| p += int(float_param) |
| assert isinstance(double_param, float) and float_equal(double_param, 1.98932, 1e-07), \ |
| 'double_param is wrong value %s ' % double_param |
| |
| p += int(double_param) |
| |
| assert boolean_param is True, 'boolean_param is wrong value %s' % boolean_param |
| |
| assert str_param == 'flink', 'str_param is wrong value %s' % str_param |
| |
| assert date_param == datetime.date(year=2014, month=9, day=13), \ |
| 'date_param is wrong value %s' % date_param |
| |
| assert time_param == datetime.time(hour=12, minute=0, second=0), \ |
| 'time_param is wrong value %s' % time_param |
| |
| assert timestamp_param == datetime.datetime(1999, 9, 10, 5, 20, 10), \ |
| 'timestamp_param is wrong value %s' % timestamp_param |
| |
| return p |
| |
| self.t_env.create_temporary_system_function("udf_with_constant_params", |
| udf(udf_with_constant_params, |
| result_type=DataTypes.BIGINT())) |
| |
| self.t_env.create_temporary_system_function( |
| "udf_with_all_constant_params", udf(lambda i, j: i + j, |
| result_type=DataTypes.BIGINT())) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b BIGINT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements([(1, 2, 3), (2, 5, 6), (3, 1, 9)], ['a', 'b', 'c']) |
| self.t_env.create_temporary_view("test_table", t) |
| self.t_env.sql_query("select udf_with_all_constant_params(" |
| "cast (1 as BIGINT)," |
| "cast (2 as BIGINT)), " |
| "udf_with_constant_params(a, " |
| "cast (null as BIGINT)," |
| "cast (1 as TINYINT)," |
| "cast (1 as SMALLINT)," |
| "cast (1 as INT)," |
| "cast (1 as BIGINT)," |
| "cast (1.05 as DECIMAL)," |
| "cast (1.23 as FLOAT)," |
| "cast (1.98932 as DOUBLE)," |
| "true," |
| "'flink'," |
| "cast ('2014-09-13' as DATE)," |
| "cast ('12:00:00' as TIME)," |
| "cast ('1999-9-10 05:20:10' as TIMESTAMP))" |
| " from test_table").execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[3, 8]", "+I[3, 9]", "+I[3, 10]"]) |
| |
| def test_overwrite_builtin_function(self): |
| self.t_env.create_temporary_system_function( |
| "plus", udf(lambda i, j: i + j - 1, |
| result_type=DataTypes.BIGINT())) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements([(1, 2, 3), (2, 5, 6), (3, 1, 9)], ['a', 'b', 'c']) |
| t.select(t.a + t.b).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[2]", "+I[6]", "+I[3]"]) |
| |
| def test_open(self): |
| self.t_env.get_config().set('python.metric.enabled', 'true') |
| execution_mode = self.t_env.get_config().get("python.execution-mode", None) |
| |
| if execution_mode == "process": |
| subtract = udf(SubtractWithMetrics(), result_type=DataTypes.BIGINT()) |
| else: |
| subtract = udf(Subtract(), result_type=DataTypes.BIGINT()) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b BIGINT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements([(1, 2), (2, 5), (3, 4)], ['a', 'b']) |
| t.select(t.a, subtract(t.b)).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[1, 1]", "+I[2, 4]", "+I[3, 3]"]) |
| |
| def test_udf_without_arguments(self): |
| one = udf(lambda: 1, result_type=DataTypes.BIGINT(), deterministic=True) |
| two = udf(lambda: 2, result_type=DataTypes.BIGINT(), deterministic=False) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a BIGINT, b BIGINT) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements([(1, 2), (2, 5), (3, 1)], ['a', 'b']) |
| t.select(one(), two()).execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[1, 2]", "+I[1, 2]", "+I[1, 2]"]) |
| |
| def test_all_data_types_expression(self): |
| @udf(result_type=DataTypes.BOOLEAN()) |
| def boolean_func(bool_param): |
| assert isinstance(bool_param, bool), 'bool_param of wrong type %s !' \ |
| % type(bool_param) |
| return bool_param |
| |
| @udf(result_type=DataTypes.TINYINT()) |
| def tinyint_func(tinyint_param): |
| assert isinstance(tinyint_param, int), 'tinyint_param of wrong type %s !' \ |
| % type(tinyint_param) |
| return tinyint_param |
| |
| @udf(result_type=DataTypes.SMALLINT()) |
| def smallint_func(smallint_param): |
| assert isinstance(smallint_param, int), 'smallint_param of wrong type %s !' \ |
| % type(smallint_param) |
| assert smallint_param == 32767, 'smallint_param of wrong value %s' % smallint_param |
| return smallint_param |
| |
| @udf(result_type=DataTypes.INT()) |
| def int_func(int_param): |
| assert isinstance(int_param, int), 'int_param of wrong type %s !' \ |
| % type(int_param) |
| assert int_param == -2147483648, 'int_param of wrong value %s' % int_param |
| return int_param |
| |
| @udf(result_type=DataTypes.BIGINT()) |
| def bigint_func(bigint_param): |
| assert isinstance(bigint_param, int), 'bigint_param of wrong type %s !' \ |
| % type(bigint_param) |
| return bigint_param |
| |
| @udf(result_type=DataTypes.BIGINT()) |
| def bigint_func_none(bigint_param): |
| assert bigint_param is None, 'bigint_param %s should be None!' % bigint_param |
| return bigint_param |
| |
| @udf(result_type=DataTypes.FLOAT()) |
| def float_func(float_param): |
| assert isinstance(float_param, float) and float_equal(float_param, 1.23, 1e-6), \ |
| 'float_param is wrong value %s !' % float_param |
| return float_param |
| |
| @udf(result_type=DataTypes.DOUBLE()) |
| def double_func(double_param): |
| assert isinstance(double_param, float) and float_equal(double_param, 1.98932, 1e-7), \ |
| 'double_param is wrong value %s !' % double_param |
| return double_param |
| |
| @udf(result_type=DataTypes.BYTES()) |
| def bytes_func(bytes_param): |
| assert bytes_param == b'flink', \ |
| 'bytes_param is wrong value %s !' % bytes_param |
| return bytes_param |
| |
| @udf(result_type=DataTypes.STRING()) |
| def str_func(str_param): |
| assert str_param == 'pyflink', \ |
| 'str_param is wrong value %s !' % str_param |
| return str_param |
| |
| @udf(result_type=DataTypes.DATE()) |
| def date_func(date_param): |
| from datetime import date |
| assert date_param == date(year=2014, month=9, day=13), \ |
| 'date_param is wrong value %s !' % date_param |
| return date_param |
| |
| @udf(result_type=DataTypes.TIME()) |
| def time_func(time_param): |
| from datetime import time |
| assert time_param == time(hour=12, minute=0, second=0, microsecond=123000), \ |
| 'time_param is wrong value %s !' % time_param |
| return time_param |
| |
| @udf(result_type=DataTypes.TIMESTAMP(3)) |
| def timestamp_func(timestamp_param): |
| from datetime import datetime |
| assert timestamp_param == datetime(2018, 3, 11, 3, 0, 0, 123000), \ |
| 'timestamp_param is wrong value %s !' % timestamp_param |
| return timestamp_param |
| |
| @udf(result_type=DataTypes.ARRAY(DataTypes.BIGINT())) |
| def array_func(array_param): |
| assert array_param == [[1, 2, 3]] or array_param == ((1, 2, 3),), \ |
| 'array_param is wrong value %s !' % array_param |
| return array_param[0] |
| |
| @udf(result_type=DataTypes.MAP(DataTypes.BIGINT(), DataTypes.STRING())) |
| def map_func(map_param): |
| assert map_param == {1: 'flink', 2: 'pyflink'}, \ |
| 'map_param is wrong value %s !' % map_param |
| return map_param |
| |
| @udf(result_type=DataTypes.DECIMAL(38, 18)) |
| def decimal_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.050000000000000000'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| @udf(result_type=DataTypes.DECIMAL(38, 18)) |
| def decimal_cut_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.059999999999999999'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| @udf(result_type=DataTypes.BINARY(5)) |
| def binary_func(binary_param): |
| assert len(binary_param) == 5 |
| return binary_param |
| |
| @udf(result_type=DataTypes.CHAR(7)) |
| def char_func(char_param): |
| assert len(char_param) == 7 |
| return char_param |
| |
| @udf(result_type=DataTypes.VARCHAR(10)) |
| def varchar_func(varchar_param): |
| assert len(varchar_param) <= 10 |
| return varchar_param |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}( |
| a BIGINT, |
| b BIGINT, |
| c TINYINT, |
| d BOOLEAN, |
| e SMALLINT, |
| f INT, |
| g FLOAT, |
| h DOUBLE, |
| i BYTES, |
| j STRING, |
| k DATE, |
| l TIME, |
| m TIMESTAMP(3), |
| n ARRAY<BIGINT>, |
| o MAP<BIGINT, STRING>, |
| p DECIMAL(38, 18), |
| q DECIMAL(38, 18), |
| r BINARY(5), |
| s CHAR(7), |
| t VARCHAR(10) |
| ) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| import datetime |
| import decimal |
| t = self.t_env.from_elements( |
| [(1, None, 1, True, 32767, -2147483648, 1.23, 1.98932, |
| bytearray(b'flink'), 'pyflink', datetime.date(2014, 9, 13), |
| datetime.time(hour=12, minute=0, second=0, microsecond=123000), |
| datetime.datetime(2018, 3, 11, 3, 0, 0, 123000), [[1, 2, 3]], |
| {1: 'flink', 2: 'pyflink'}, decimal.Decimal('1000000000000000000.05'), |
| decimal.Decimal('1000000000000000000.05999999999999999899999999999'), |
| bytearray(b'flink'), 'pyflink', 'pyflink')], |
| DataTypes.ROW( |
| [DataTypes.FIELD("a", DataTypes.BIGINT()), |
| DataTypes.FIELD("b", DataTypes.BIGINT()), |
| DataTypes.FIELD("c", DataTypes.TINYINT()), |
| DataTypes.FIELD("d", DataTypes.BOOLEAN()), |
| DataTypes.FIELD("e", DataTypes.SMALLINT()), |
| DataTypes.FIELD("f", DataTypes.INT()), |
| DataTypes.FIELD("g", DataTypes.FLOAT()), |
| DataTypes.FIELD("h", DataTypes.DOUBLE()), |
| DataTypes.FIELD("i", DataTypes.BYTES()), |
| DataTypes.FIELD("j", DataTypes.STRING()), |
| DataTypes.FIELD("k", DataTypes.DATE()), |
| DataTypes.FIELD("l", DataTypes.TIME()), |
| DataTypes.FIELD("m", DataTypes.TIMESTAMP(3)), |
| DataTypes.FIELD("n", DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.BIGINT()))), |
| DataTypes.FIELD("o", DataTypes.MAP(DataTypes.BIGINT(), DataTypes.STRING())), |
| DataTypes.FIELD("p", DataTypes.DECIMAL(38, 18)), |
| DataTypes.FIELD("q", DataTypes.DECIMAL(38, 18)), |
| DataTypes.FIELD("r", DataTypes.BINARY(5)), |
| DataTypes.FIELD("s", DataTypes.CHAR(7)), |
| DataTypes.FIELD("t", DataTypes.VARCHAR(10))])) |
| |
| t.select( |
| bigint_func(t.a), |
| bigint_func_none(t.b), |
| tinyint_func(t.c), |
| boolean_func(t.d), |
| smallint_func(t.e), |
| int_func(t.f), |
| float_func(t.g), |
| double_func(t.h), |
| bytes_func(t.i), |
| str_func(t.j), |
| date_func(t.k), |
| time_func(t.l), |
| timestamp_func(t.m), |
| array_func(t.n), |
| map_func(t.o), |
| decimal_func(t.p), |
| decimal_cut_func(t.q), |
| binary_func(t.r), |
| char_func(t.s), |
| varchar_func(t.t)) \ |
| .execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| # Currently the sink result precision of DataTypes.TIME(precision) only supports 0. |
| self.assert_equals(actual, |
| ["+I[1, null, 1, true, 32767, -2147483648, 1.23, 1.98932, " |
| "[102, 108, 105, 110, 107], pyflink, 2014-09-13, 12:00:00.123, " |
| "2018-03-11T03:00:00.123, [1, 2, 3], " |
| "{1=flink, 2=pyflink}, 1000000000000000000.050000000000000000, " |
| "1000000000000000000.059999999999999999, [102, 108, 105, 110, 107], " |
| "pyflink, pyflink]"]) |
| |
| def test_all_data_types(self): |
| def boolean_func(bool_param): |
| assert isinstance(bool_param, bool), 'bool_param of wrong type %s !' \ |
| % type(bool_param) |
| return bool_param |
| |
| def tinyint_func(tinyint_param): |
| assert isinstance(tinyint_param, int), 'tinyint_param of wrong type %s !' \ |
| % type(tinyint_param) |
| return tinyint_param |
| |
| def smallint_func(smallint_param): |
| assert isinstance(smallint_param, int), 'smallint_param of wrong type %s !' \ |
| % type(smallint_param) |
| assert smallint_param == 32767, 'smallint_param of wrong value %s' % smallint_param |
| return smallint_param |
| |
| def int_func(int_param): |
| assert isinstance(int_param, int), 'int_param of wrong type %s !' \ |
| % type(int_param) |
| assert int_param == -2147483648, 'int_param of wrong value %s' % int_param |
| return int_param |
| |
| def bigint_func(bigint_param): |
| assert isinstance(bigint_param, int), 'bigint_param of wrong type %s !' \ |
| % type(bigint_param) |
| return bigint_param |
| |
| def bigint_func_none(bigint_param): |
| assert bigint_param is None, 'bigint_param %s should be None!' % bigint_param |
| return bigint_param |
| |
| def float_func(float_param): |
| assert isinstance(float_param, float) and float_equal(float_param, 1.23, 1e-6), \ |
| 'float_param is wrong value %s !' % float_param |
| return float_param |
| |
| def double_func(double_param): |
| assert isinstance(double_param, float) and float_equal(double_param, 1.98932, 1e-7), \ |
| 'double_param is wrong value %s !' % double_param |
| return double_param |
| |
| def bytes_func(bytes_param): |
| assert bytes_param == b'flink', \ |
| 'bytes_param is wrong value %s !' % bytes_param |
| return bytes_param |
| |
| def str_func(str_param): |
| assert str_param == 'pyflink', \ |
| 'str_param is wrong value %s !' % str_param |
| return str_param |
| |
| def date_func(date_param): |
| from datetime import date |
| assert date_param == date(year=2014, month=9, day=13), \ |
| 'date_param is wrong value %s !' % date_param |
| return date_param |
| |
| def time_func(time_param): |
| from datetime import time |
| assert time_param == time(hour=12, minute=0, second=0, microsecond=123000), \ |
| 'time_param is wrong value %s !' % time_param |
| return time_param |
| |
| def timestamp_func(timestamp_param): |
| from datetime import datetime |
| assert timestamp_param == datetime(2018, 3, 11, 3, 0, 0, 123000), \ |
| 'timestamp_param is wrong value %s !' % timestamp_param |
| return timestamp_param |
| |
| def array_func(array_param): |
| assert array_param == [[1, 2, 3]] or array_param == ((1, 2, 3),), \ |
| 'array_param is wrong value %s !' % array_param |
| return array_param[0] |
| |
| def map_func(map_param): |
| assert map_param == {1: 'flink', 2: 'pyflink'}, \ |
| 'map_param is wrong value %s !' % map_param |
| return map_param |
| |
| def decimal_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.050000000000000000'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| def decimal_cut_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.059999999999999999'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| self.t_env.create_temporary_system_function( |
| "boolean_func", udf(boolean_func, result_type=DataTypes.BOOLEAN())) |
| |
| self.t_env.create_temporary_system_function( |
| "tinyint_func", udf(tinyint_func, result_type=DataTypes.TINYINT())) |
| |
| self.t_env.create_temporary_system_function( |
| "smallint_func", udf(smallint_func, result_type=DataTypes.SMALLINT())) |
| |
| self.t_env.create_temporary_system_function( |
| "int_func", udf(int_func, result_type=DataTypes.INT())) |
| |
| self.t_env.create_temporary_system_function( |
| "bigint_func", udf(bigint_func, result_type=DataTypes.BIGINT())) |
| |
| self.t_env.create_temporary_system_function( |
| "bigint_func_none", udf(bigint_func_none, result_type=DataTypes.BIGINT())) |
| |
| self.t_env.create_temporary_system_function( |
| "float_func", udf(float_func, result_type=DataTypes.FLOAT())) |
| |
| self.t_env.create_temporary_system_function( |
| "double_func", udf(double_func, result_type=DataTypes.DOUBLE())) |
| |
| self.t_env.create_temporary_system_function( |
| "bytes_func", udf(bytes_func, result_type=DataTypes.BYTES())) |
| |
| self.t_env.create_temporary_system_function( |
| "str_func", udf(str_func, result_type=DataTypes.STRING())) |
| |
| self.t_env.create_temporary_system_function( |
| "date_func", udf(date_func, result_type=DataTypes.DATE())) |
| |
| self.t_env.create_temporary_system_function( |
| "time_func", udf(time_func, result_type=DataTypes.TIME())) |
| |
| self.t_env.create_temporary_system_function( |
| "timestamp_func", udf(timestamp_func, result_type=DataTypes.TIMESTAMP(3))) |
| |
| self.t_env.create_temporary_system_function( |
| "array_func", udf(array_func, result_type=DataTypes.ARRAY(DataTypes.BIGINT()))) |
| |
| self.t_env.create_temporary_system_function( |
| "map_func", udf(map_func, |
| result_type=DataTypes.MAP(DataTypes.BIGINT(), DataTypes.STRING()))) |
| |
| self.t_env.create_temporary_system_function( |
| "decimal_func", udf(decimal_func, result_type=DataTypes.DECIMAL(38, 18))) |
| |
| self.t_env.create_temporary_system_function( |
| "decimal_cut_func", udf(decimal_cut_func, result_type=DataTypes.DECIMAL(38, 18))) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}( |
| a BIGINT, b BIGINT, c TINYINT, d BOOLEAN, e SMALLINT, f INT, g FLOAT, h DOUBLE, |
| i BYTES, j STRING, k DATE, l TIME, m TIMESTAMP(3), n ARRAY<BIGINT>, |
| o MAP<BIGINT, STRING>, p DECIMAL(38, 18), q DECIMAL(38, 18)) |
| WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| import datetime |
| import decimal |
| t = self.t_env.from_elements( |
| [(1, None, 1, True, 32767, -2147483648, 1.23, 1.98932, |
| bytearray(b'flink'), 'pyflink', datetime.date(2014, 9, 13), |
| datetime.time(hour=12, minute=0, second=0, microsecond=123000), |
| datetime.datetime(2018, 3, 11, 3, 0, 0, 123000), [[1, 2, 3]], |
| {1: 'flink', 2: 'pyflink'}, decimal.Decimal('1000000000000000000.05'), |
| decimal.Decimal('1000000000000000000.05999999999999999899999999999'))], |
| DataTypes.ROW( |
| [DataTypes.FIELD("a", DataTypes.BIGINT()), |
| DataTypes.FIELD("b", DataTypes.BIGINT()), |
| DataTypes.FIELD("c", DataTypes.TINYINT()), |
| DataTypes.FIELD("d", DataTypes.BOOLEAN()), |
| DataTypes.FIELD("e", DataTypes.SMALLINT()), |
| DataTypes.FIELD("f", DataTypes.INT()), |
| DataTypes.FIELD("g", DataTypes.FLOAT()), |
| DataTypes.FIELD("h", DataTypes.DOUBLE()), |
| DataTypes.FIELD("i", DataTypes.BYTES()), |
| DataTypes.FIELD("j", DataTypes.STRING()), |
| DataTypes.FIELD("k", DataTypes.DATE()), |
| DataTypes.FIELD("l", DataTypes.TIME()), |
| DataTypes.FIELD("m", DataTypes.TIMESTAMP(3)), |
| DataTypes.FIELD("n", DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.BIGINT()))), |
| DataTypes.FIELD("o", DataTypes.MAP(DataTypes.BIGINT(), DataTypes.STRING())), |
| DataTypes.FIELD("p", DataTypes.DECIMAL(38, 18)), |
| DataTypes.FIELD("q", DataTypes.DECIMAL(38, 18))])) |
| |
| t.select(call("bigint_func", t.a), call("bigint_func_none", t.b), |
| call("tinyint_func", t.c), call("boolean_func", t.d), |
| call("smallint_func", t.e), call("int_func", t.f), |
| call("float_func", t.g), call("double_func", t.h), |
| call("bytes_func", t.i), call("str_func", t.j), |
| call("date_func", t.k), call("time_func", t.l), |
| call("timestamp_func", t.m), call("array_func", t.n), |
| call("map_func", t.o), call("decimal_func", t.p), |
| call("decimal_cut_func", t.q)) \ |
| .execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| # Currently the sink result precision of DataTypes.TIME(precision) only supports 0. |
| self.assert_equals(actual, |
| ["+I[1, null, 1, true, 32767, -2147483648, 1.23, 1.98932, " |
| "[102, 108, 105, 110, 107], pyflink, 2014-09-13, " |
| "12:00:00.123, 2018-03-11T03:00:00.123, [1, 2, 3], " |
| "{1=flink, 2=pyflink}, 1000000000000000000.050000000000000000, " |
| "1000000000000000000.059999999999999999]"]) |
| |
| def test_using_udf_in_sql(self): |
| sql = f""" |
| CREATE TEMPORARY FUNCTION echo AS |
| '{UserDefinedFunctionTests.__module__}.echo' |
| LANGUAGE PYTHON |
| """ |
| self.t_env.execute_sql(sql) |
| self.assert_equals( |
| list(self.t_env.execute_sql("SELECT echo(1)").collect()), |
| [Row(1)]) |
| |
| def test_all_data_types_string(self): |
| @udf(result_type='BOOLEAN') |
| def boolean_func(bool_param): |
| assert isinstance(bool_param, bool), 'bool_param of wrong type %s !' \ |
| % type(bool_param) |
| return bool_param |
| |
| @udf(result_type='TINYINT') |
| def tinyint_func(tinyint_param): |
| assert isinstance(tinyint_param, int), 'tinyint_param of wrong type %s !' \ |
| % type(tinyint_param) |
| return tinyint_param |
| |
| @udf(result_type='SMALLINT') |
| def smallint_func(smallint_param): |
| assert isinstance(smallint_param, int), 'smallint_param of wrong type %s !' \ |
| % type(smallint_param) |
| assert smallint_param == 32767, 'smallint_param of wrong value %s' % smallint_param |
| return smallint_param |
| |
| @udf(result_type='INT') |
| def int_func(int_param): |
| assert isinstance(int_param, int), 'int_param of wrong type %s !' \ |
| % type(int_param) |
| assert int_param == -2147483648, 'int_param of wrong value %s' % int_param |
| return int_param |
| |
| @udf(result_type='BIGINT') |
| def bigint_func(bigint_param): |
| assert isinstance(bigint_param, int), 'bigint_param of wrong type %s !' \ |
| % type(bigint_param) |
| return bigint_param |
| |
| @udf(result_type='BIGINT') |
| def bigint_func_none(bigint_param): |
| assert bigint_param is None, 'bigint_param %s should be None!' % bigint_param |
| return bigint_param |
| |
| @udf(result_type='FLOAT') |
| def float_func(float_param): |
| assert isinstance(float_param, float) and float_equal(float_param, 1.23, 1e-6), \ |
| 'float_param is wrong value %s !' % float_param |
| return float_param |
| |
| @udf(result_type='DOUBLE') |
| def double_func(double_param): |
| assert isinstance(double_param, float) and float_equal(double_param, 1.98932, 1e-7), \ |
| 'double_param is wrong value %s !' % double_param |
| return double_param |
| |
| @udf(result_type='BYTES') |
| def bytes_func(bytes_param): |
| assert bytes_param == b'flink', \ |
| 'bytes_param is wrong value %s !' % bytes_param |
| return bytes_param |
| |
| @udf(result_type='STRING') |
| def str_func(str_param): |
| assert str_param == 'pyflink', \ |
| 'str_param is wrong value %s !' % str_param |
| return str_param |
| |
| @udf(result_type='DATE') |
| def date_func(date_param): |
| from datetime import date |
| assert date_param == date(year=2014, month=9, day=13), \ |
| 'date_param is wrong value %s !' % date_param |
| return date_param |
| |
| @udf(result_type='TIME') |
| def time_func(time_param): |
| from datetime import time |
| assert time_param == time(hour=12, minute=0, second=0, microsecond=123000), \ |
| 'time_param is wrong value %s !' % time_param |
| return time_param |
| |
| @udf(result_type='TIMESTAMP(3)') |
| def timestamp_func(timestamp_param): |
| from datetime import datetime |
| assert timestamp_param == datetime(2018, 3, 11, 3, 0, 0, 123000), \ |
| 'timestamp_param is wrong value %s !' % timestamp_param |
| return timestamp_param |
| |
| @udf(result_type='ARRAY<BIGINT>') |
| def array_func(array_param): |
| assert array_param == [[1, 2, 3]] or array_param == ((1, 2, 3),), \ |
| 'array_param is wrong value %s !' % array_param |
| return array_param[0] |
| |
| @udf(result_type='MAP<BIGINT, STRING>') |
| def map_func(map_param): |
| assert map_param == {1: 'flink', 2: 'pyflink'}, \ |
| 'map_param is wrong value %s !' % map_param |
| return map_param |
| |
| @udf(result_type='DECIMAL(38, 18)') |
| def decimal_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.050000000000000000'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| @udf(result_type='DECIMAL(38, 18)') |
| def decimal_cut_func(decimal_param): |
| from decimal import Decimal |
| assert decimal_param == Decimal('1000000000000000000.059999999999999999'), \ |
| 'decimal_param is wrong value %s !' % decimal_param |
| return decimal_param |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}( |
| a BIGINT, b BIGINT, c TINYINT, d BOOLEAN, e SMALLINT, f INT, g FLOAT, h DOUBLE, |
| i BYTES, j STRING, k DATE, l TIME, m TIMESTAMP(3), n ARRAY<BIGINT>, |
| o MAP<BIGINT, STRING>, p DECIMAL(38, 18), q DECIMAL(38, 18)) |
| WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| import datetime |
| import decimal |
| t = self.t_env.from_elements( |
| [(1, None, 1, True, 32767, -2147483648, 1.23, 1.98932, |
| bytearray(b'flink'), 'pyflink', datetime.date(2014, 9, 13), |
| datetime.time(hour=12, minute=0, second=0, microsecond=123000), |
| datetime.datetime(2018, 3, 11, 3, 0, 0, 123000), [[1, 2, 3]], |
| {1: 'flink', 2: 'pyflink'}, decimal.Decimal('1000000000000000000.05'), |
| decimal.Decimal('1000000000000000000.05999999999999999899999999999'))], |
| DataTypes.ROW( |
| [DataTypes.FIELD("a", DataTypes.BIGINT()), |
| DataTypes.FIELD("b", DataTypes.BIGINT()), |
| DataTypes.FIELD("c", DataTypes.TINYINT()), |
| DataTypes.FIELD("d", DataTypes.BOOLEAN()), |
| DataTypes.FIELD("e", DataTypes.SMALLINT()), |
| DataTypes.FIELD("f", DataTypes.INT()), |
| DataTypes.FIELD("g", DataTypes.FLOAT()), |
| DataTypes.FIELD("h", DataTypes.DOUBLE()), |
| DataTypes.FIELD("i", DataTypes.BYTES()), |
| DataTypes.FIELD("j", DataTypes.STRING()), |
| DataTypes.FIELD("k", DataTypes.DATE()), |
| DataTypes.FIELD("l", DataTypes.TIME()), |
| DataTypes.FIELD("m", DataTypes.TIMESTAMP(3)), |
| DataTypes.FIELD("n", DataTypes.ARRAY(DataTypes.ARRAY(DataTypes.BIGINT()))), |
| DataTypes.FIELD("o", DataTypes.MAP(DataTypes.BIGINT(), DataTypes.STRING())), |
| DataTypes.FIELD("p", DataTypes.DECIMAL(38, 18)), |
| DataTypes.FIELD("q", DataTypes.DECIMAL(38, 18))])) |
| |
| t.select( |
| bigint_func(t.a), |
| bigint_func_none(t.b), |
| tinyint_func(t.c), |
| boolean_func(t.d), |
| smallint_func(t.e), |
| int_func(t.f), |
| float_func(t.g), |
| double_func(t.h), |
| bytes_func(t.i), |
| str_func(t.j), |
| date_func(t.k), |
| time_func(t.l), |
| timestamp_func(t.m), |
| array_func(t.n), |
| map_func(t.o), |
| decimal_func(t.p), |
| decimal_cut_func(t.q)) \ |
| .execute_insert(sink_table).wait() |
| actual = source_sink_utils.results() |
| # Currently the sink result precision of DataTypes.TIME(precision) only supports 0. |
| self.assert_equals(actual, |
| ["+I[1, null, 1, true, 32767, -2147483648, 1.23, 1.98932, " |
| "[102, 108, 105, 110, 107], pyflink, 2014-09-13, 12:00:00.123, " |
| "2018-03-11T03:00:00.123, [1, 2, 3], " |
| "{1=flink, 2=pyflink}, 1000000000000000000.050000000000000000, " |
| "1000000000000000000.059999999999999999]"]) |
| |
| def test_create_and_drop_function(self): |
| t_env = self.t_env |
| |
| t_env.create_temporary_system_function( |
| "add_one_func", udf(lambda i: i + 1, result_type=DataTypes.BIGINT())) |
| t_env.create_temporary_function( |
| "subtract_one_func", udf(SubtractOne(), result_type=DataTypes.BIGINT())) |
| self.assertTrue('add_one_func' in t_env.list_user_defined_functions()) |
| self.assertTrue('subtract_one_func' in t_env.list_user_defined_functions()) |
| |
| t_env.drop_temporary_system_function("add_one_func") |
| t_env.drop_temporary_function("subtract_one_func") |
| self.assertTrue('add_one_func' not in t_env.list_user_defined_functions()) |
| self.assertTrue('subtract_one_func' not in t_env.list_user_defined_functions()) |
| |
| |
| # decide whether two floats are equal |
| def float_equal(a, b, rel_tol=1e-09, abs_tol=0.0): |
| return abs(a - b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol) |
| |
| |
| class PyFlinkStreamUserDefinedFunctionTests(UserDefinedFunctionTests, |
| PyFlinkStreamTableTestCase): |
| def test_deterministic(self): |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) |
| self.assertTrue(add_one._deterministic) |
| |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT(), deterministic=False) |
| self.assertFalse(add_one._deterministic) |
| |
| subtract_one = udf(SubtractOne(), result_type=DataTypes.BIGINT()) |
| self.assertTrue(subtract_one._deterministic) |
| |
| with self.assertRaises(ValueError, msg="Inconsistent deterministic: False and True"): |
| udf(SubtractOne(), result_type=DataTypes.BIGINT(), deterministic=False) |
| |
| self.assertTrue(add._deterministic) |
| |
| @udf(result_type=DataTypes.BIGINT(), deterministic=False) |
| def non_deterministic_udf(i): |
| return i |
| |
| self.assertFalse(non_deterministic_udf._deterministic) |
| |
| def test_name(self): |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) |
| self.assertEqual("<lambda>", add_one._name) |
| |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT(), name="add_one") |
| self.assertEqual("add_one", add_one._name) |
| |
| subtract_one = udf(SubtractOne(), result_type=DataTypes.BIGINT()) |
| self.assertEqual("SubtractOne", subtract_one._name) |
| |
| subtract_one = udf(SubtractOne(), result_type=DataTypes.BIGINT(), name="subtract_one") |
| self.assertEqual("subtract_one", subtract_one._name) |
| |
| self.assertEqual("add", add._name) |
| |
| @udf(result_type=DataTypes.BIGINT(), name="named") |
| def named_udf(i): |
| return i |
| |
| self.assertEqual("named", named_udf._name) |
| |
| def test_abc(self): |
| class UdfWithoutEval(ScalarFunction): |
| def open(self, function_context): |
| pass |
| |
| with self.assertRaises( |
| TypeError, |
| msg="Can't instantiate abstract class UdfWithoutEval with abstract methods eval"): |
| UdfWithoutEval() |
| |
| def test_invalid_udf(self): |
| class Plus(object): |
| def eval(self, col): |
| return col + 1 |
| |
| with self.assertRaises( |
| TypeError, |
| msg="Invalid function: not a function or callable (__call__ is not defined)"): |
| # test non-callable function |
| self.t_env.create_temporary_system_function( |
| "non-callable-udf", udf(Plus(), DataTypes.BIGINT(), DataTypes.BIGINT())) |
| |
| def test_data_types(self): |
| timezone = self.t_env.get_config().get_local_timezone() |
| local_datetime = pytz.timezone(timezone).localize( |
| datetime.datetime(1970, 1, 1, 0, 0, 0, 123000)) |
| |
| @udf(result_type=DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3)) |
| def local_zoned_timestamp_func(local_zoned_timestamp_param): |
| assert local_zoned_timestamp_param == local_datetime, \ |
| 'local_zoned_timestamp_param is wrong value %s !' % local_zoned_timestamp_param |
| return local_zoned_timestamp_param |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}(a TIMESTAMP_LTZ(3)) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| |
| t = self.t_env.from_elements( |
| [(local_datetime,)], |
| DataTypes.ROW([DataTypes.FIELD("a", DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(3))])) |
| |
| t.select(local_zoned_timestamp_func(local_zoned_timestamp_func(t.a))) \ |
| .execute_insert(sink_table) \ |
| .wait() |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ["+I[1970-01-01T00:00:00.123Z]"]) |
| |
| def test_execute_from_json_plan(self): |
| # create source file path |
| tmp_dir = self.tempdir |
| data = ['1,1', '3,3', '2,2'] |
| source_path = tmp_dir + '/test_execute_from_json_plan_input.csv' |
| sink_path = tmp_dir + '/test_execute_from_json_plan_out' |
| with open(source_path, 'w') as fd: |
| for ele in data: |
| fd.write(ele + '\n') |
| |
| source_table = """ |
| CREATE TABLE source_table ( |
| a BIGINT, |
| b BIGINT |
| ) WITH ( |
| 'connector' = 'filesystem', |
| 'path' = '%s', |
| 'format' = 'csv' |
| ) |
| """ % source_path |
| self.t_env.execute_sql(source_table) |
| |
| self.t_env.execute_sql(""" |
| CREATE TABLE sink_table ( |
| id BIGINT, |
| data BIGINT |
| ) WITH ( |
| 'connector' = 'filesystem', |
| 'path' = '%s', |
| 'format' = 'csv' |
| ) |
| """ % sink_path) |
| |
| add_one = udf(lambda i: i + 1, result_type=DataTypes.BIGINT()) |
| self.t_env.create_temporary_system_function("add_one", add_one) |
| |
| json_plan = self.t_env._j_tenv.compilePlanSql("INSERT INTO sink_table SELECT " |
| "a, " |
| "add_one(b) " |
| "FROM source_table") |
| from py4j.java_gateway import get_method |
| get_method(json_plan.execute(), "await")() |
| |
| import glob |
| lines = [line.strip() for file in glob.glob(sink_path + '/*') for line in open(file, 'r')] |
| lines.sort() |
| self.assertEqual(lines, ['1,2', '2,3', '3,4']) |
| |
| def test_udf_with_rowtime_arguments(self): |
| from pyflink.common import WatermarkStrategy |
| from pyflink.common.typeinfo import Types |
| from pyflink.common.watermark_strategy import TimestampAssigner |
| from pyflink.table import Schema |
| |
| class MyTimestampAssigner(TimestampAssigner): |
| |
| def extract_timestamp(self, value, record_timestamp) -> int: |
| return int(value[0]) |
| |
| ds = self.env.from_collection( |
| [(1, 42, "a"), (2, 5, "a"), (3, 1000, "c"), (100, 1000, "c")], |
| Types.ROW_NAMED(["a", "b", "c"], [Types.LONG(), Types.INT(), Types.STRING()])) |
| |
| ds = ds.assign_timestamps_and_watermarks( |
| WatermarkStrategy.for_monotonous_timestamps() |
| .with_timestamp_assigner(MyTimestampAssigner())) |
| |
| table = self.t_env.from_data_stream( |
| ds, |
| Schema.new_builder() |
| .column_by_metadata("rowtime", "TIMESTAMP_LTZ(3)") |
| .watermark("rowtime", "SOURCE_WATERMARK()") |
| .build()) |
| |
| @udf(result_type=DataTypes.ROW([DataTypes.FIELD('f1', DataTypes.INT())])) |
| def inc(input_row): |
| return Row(input_row.b) |
| |
| sink_table = generate_random_table_name() |
| sink_table_ddl = f""" |
| CREATE TABLE {sink_table}( |
| a INT |
| ) WITH ('connector'='test-sink') |
| """ |
| self.t_env.execute_sql(sink_table_ddl) |
| table.map(inc).execute_insert(sink_table).wait() |
| |
| actual = source_sink_utils.results() |
| self.assert_equals(actual, ['+I[42]', '+I[5]', '+I[1000]', '+I[1000]']) |
| |
| |
| class PyFlinkBatchUserDefinedFunctionTests(UserDefinedFunctionTests, |
| PyFlinkBatchTableTestCase): |
| pass |
| |
| |
| class PyFlinkEmbeddedThreadTests(UserDefinedFunctionTests, PyFlinkBatchTableTestCase): |
| def setUp(self): |
| super(PyFlinkEmbeddedThreadTests, self).setUp() |
| self.t_env.get_config().set("python.execution-mode", "thread") |
| |
| |
| # test specify the input_types |
| @udf(input_types=[DataTypes.BIGINT(), DataTypes.BIGINT()], result_type=DataTypes.BIGINT()) |
| def add(i, j): |
| return i + j |
| |
| |
| class SubtractOne(ScalarFunction): |
| |
| def eval(self, i): |
| return i - 1 |
| |
| |
| class SubtractWithParameters(ScalarFunction): |
| |
| def open(self, function_context: FunctionContext): |
| self.subtract_value = int(function_context.get_job_parameter("subtract_value", "1")) |
| |
| def eval(self, i): |
| return i - self.subtract_value |
| |
| |
| class SubtractWithMetrics(ScalarFunction, unittest.TestCase): |
| |
| def open(self, function_context): |
| self.subtracted_value = 1 |
| mg = function_context.get_metric_group() |
| self.counter = mg.add_group("key", "value").counter("my_counter") |
| self.counter_sum = 0 |
| |
| def eval(self, i): |
| # counter |
| self.counter.inc(i) |
| self.counter_sum += i |
| return i - self.subtracted_value |
| |
| |
| class Subtract(ScalarFunction, unittest.TestCase): |
| |
| def open(self, function_context): |
| self.subtracted_value = 1 |
| self.counter_sum = 0 |
| |
| def eval(self, i): |
| # counter |
| self.counter_sum += i |
| return i - self.subtracted_value |
| |
| |
| class CallablePlus(object): |
| |
| def __call__(self, col): |
| return col + 1 |
| |
| |
| @udf(input_types=['BIGINT'], result_type='BIGINT') |
| def echo(i: str): |
| return i |
| |
| |
| if __name__ == '__main__': |
| import unittest |
| |
| try: |
| import xmlrunner |
| |
| testRunner = xmlrunner.XMLTestRunner(output='target/test-reports') |
| except ImportError: |
| testRunner = None |
| unittest.main(testRunner=testRunner, verbosity=2) |
