/* * 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. */ package org.apache.flink.table.runtime.runners.python.table; import org.apache.flink.api.common.typeutils.TypeSerializer; import org.apache.flink.api.java.typeutils.runtime.RowSerializer; import org.apache.flink.fnexecution.v1.FlinkFnApi; import org.apache.flink.python.env.PythonEnvironmentManager; import org.apache.flink.python.metric.FlinkMetricContainer; import org.apache.flink.table.functions.python.PythonFunctionInfo; import org.apache.flink.table.runtime.runners.python.scalar.AbstractPythonScalarFunctionRunnerTest; import org.apache.flink.table.runtime.utils.PythonTestUtils; import org.apache.flink.table.types.logical.BigIntType; import org.apache.flink.table.types.logical.RowType; import org.apache.flink.types.Row; import org.apache.beam.runners.fnexecution.control.JobBundleFactory; import org.apache.beam.sdk.fn.data.FnDataReceiver; import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct; import org.junit.Test; import java.io.IOException; import java.util.Collections; import static org.apache.flink.table.runtime.utils.PythonTestUtils.createTestEnvironmentManager; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; /** * Tests for {@link PythonTableFunctionRunner}. These test that: * * <ul>T * <li>The input data type and output data type are properly constructed</li> * <li>The UDTF proto is properly constructed</li> * </ul> */ public class PythonTableFunctionRunnerTest extends AbstractPythonTableFunctionRunnerTest<Row> { @Test public void testInputOutputDataTypeConstructedProperlyForSingleUDTF() throws Exception { final AbstractPythonTableFunctionRunner<Row> runner = createUDTFRunner(); // check input TypeSerializer TypeSerializer inputTypeSerializer = runner.getInputTypeSerializer(); assertTrue(inputTypeSerializer instanceof RowSerializer); assertEquals(1, ((RowSerializer) inputTypeSerializer).getArity()); } @Test public void testUDFnProtoConstructedProperlyForSingleUTDF() throws Exception { final AbstractPythonTableFunctionRunner<Row> runner = createUDTFRunner(); FlinkFnApi.UserDefinedFunctions udtfs = runner.getUserDefinedFunctionsProto(); assertEquals(1, udtfs.getUdfsCount()); FlinkFnApi.UserDefinedFunction udtf = udtfs.getUdfs(0); assertEquals(1, udtf.getInputsCount()); assertEquals(0, udtf.getInputs(0).getInputOffset()); } @Override public AbstractPythonTableFunctionRunner<Row> createPythonTableFunctionRunner( PythonFunctionInfo pythonFunctionInfo, RowType inputType, RowType outputType) throws Exception { final FnDataReceiver<byte[]> dummyReceiver = input -> { // ignore the execution results }; final PythonEnvironmentManager environmentManager = createTestEnvironmentManager(); return new PythonTableFunctionRunner( "testPythonRunner", dummyReceiver, pythonFunctionInfo, environmentManager, inputType, outputType, Collections.emptyMap(), PythonTestUtils.createMockFlinkMetricContainer()); } private AbstractPythonTableFunctionRunner<Row> createUDTFRunner( JobBundleFactory jobBundleFactory, FnDataReceiver<byte[]> receiver) throws IOException { PythonFunctionInfo pythonFunctionInfo = new PythonFunctionInfo( AbstractPythonScalarFunctionRunnerTest.DummyPythonFunction.INSTANCE, new Integer[]{0}); RowType rowType = new RowType(Collections.singletonList(new RowType.RowField("f1", new BigIntType()))); final PythonEnvironmentManager environmentManager = createTestEnvironmentManager(); return new PythonTableFunctionRunnerTestHarness( "testPythonRunner", receiver, pythonFunctionInfo, environmentManager, rowType, rowType, jobBundleFactory, PythonTestUtils.createMockFlinkMetricContainer()); } private static class PythonTableFunctionRunnerTestHarness extends PythonTableFunctionRunner { private final JobBundleFactory jobBundleFactory; PythonTableFunctionRunnerTestHarness( String taskName, FnDataReceiver<byte[]> resultReceiver, PythonFunctionInfo tableFunction, PythonEnvironmentManager environmentManager, RowType inputType, RowType outputType, JobBundleFactory jobBundleFactory, FlinkMetricContainer flinkMetricContainer) { super(taskName, resultReceiver, tableFunction, environmentManager, inputType, outputType, Collections.emptyMap(), flinkMetricContainer); this.jobBundleFactory = jobBundleFactory; } @Override public JobBundleFactory createJobBundleFactory(Struct pipelineOptions) throws Exception { return jobBundleFactory; } } }