def visit_create_udf(self, statement: CreateUDFStatement):
"""Convertor for parsed create udf statement
Arguments:
statement {CreateUDFStatement} -- Create UDF Statement
"""
annotated_inputs = column_definition_to_udf_io(statement.inputs, True)
annotated_outputs = column_definition_to_udf_io(
statement.outputs, False)
create_udf_opr = LogicalCreateUDF(statement.name,
statement.if_not_exists,
annotated_inputs, annotated_outputs,
statement.impl_path,
statement.udf_type)
self._plan = create_udf_opr
def test_column_definition_to_udf_io(self, mock_func, mock):
mock.return_value.udf_io.return_value = 'udf_io'
col = MagicMock(spec=ColumnDefinition)
col.name = 'name'
col.type = 'type'
col.dimension = 'dimension'
col_list = [col, col]
mock_func.return_value = col.type
actual = column_definition_to_udf_io(col_list, True)
for col in col_list:
mock_func.assert_called_with('type')
mock.return_value.udf_io.assert_called_with(
'name', 'type', 'dimension', True)
self.assertEqual(actual, ['udf_io', 'udf_io'])
# input not list
actual2 = column_definition_to_udf_io(col, True)
mock.return_value.udf_io.assert_called_with('name', 'type',
'dimension', True)
self.assertEqual(actual2, ['udf_io'])
def test_column_definition_to_udf_io(self):
col = ColumnDefinition('data', ColumnType.NDARRAY, NdArrayType.UINT8,
[None, None, None])
col_list = [col, col]
actual = column_definition_to_udf_io(col_list, True)
for io in actual:
self.assertEqual(io.name, 'data')
self.assertEqual(io.type, ColumnType.NDARRAY)
self.assertEqual(io.is_nullable, False)
self.assertEqual(io.array_type, NdArrayType.UINT8)
self.assertEqual(io.array_dimensions, [None, None, None])
self.assertEqual(io.is_input, True)
self.assertEqual(io.udf_id, None)
# input not list
actual2 = column_definition_to_udf_io(col, True)
for io in actual2:
self.assertEqual(io.name, 'data')
self.assertEqual(io.type, ColumnType.NDARRAY)
self.assertEqual(io.is_nullable, False)
self.assertEqual(io.array_type, NdArrayType.UINT8)
self.assertEqual(io.array_dimensions, [None, None, None])
self.assertEqual(io.is_input, True)
self.assertEqual(io.udf_id, None)
