You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by he...@apache.org on 2019/09/23 12:12:14 UTC
[flink] branch master updated: [FLINK-14015][python] Introduces
PythonScalarFunctionOperator to execute Python user-defined functions
This is an automated email from the ASF dual-hosted git repository.
hequn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 71fa237 [FLINK-14015][python] Introduces PythonScalarFunctionOperator to execute Python user-defined functions
71fa237 is described below
commit 71fa23738e7f5f582c20753a92653bc3ce1e29b8
Author: Dian Fu <fu...@alibaba-inc.com>
AuthorDate: Tue Sep 17 18:09:00 2019 +0800
[FLINK-14015][python] Introduces PythonScalarFunctionOperator to execute Python user-defined functions
- Introduces PythonScalarFunctionOperator to execute Python user-defined functions for flink planner.
- Introduces BaseRowPythonScalarFunctionOperator to execute Python user-defined functions for blink planner.
This closes #9707.
---
.../src/main/java/org/apache/flink/types/Row.java | 30 +++
.../test/java/org/apache/flink/types/RowTest.java | 24 ++
flink-python/pom.xml | 6 +
.../flink/python/AbstractPythonFunctionRunner.java | 34 +--
.../org/apache/flink/python/PythonOptions.java | 48 ++++
.../python/AbstractPythonFunctionOperator.java | 278 +++++++++++++++++++++
.../AbstractPythonScalarFunctionOperator.java | 200 +++++++++++++++
.../BaseRowPythonScalarFunctionOperator.java | 178 +++++++++++++
.../python/PythonScalarFunctionOperator.java | 158 ++++++++++++
.../python/AbstractPythonScalarFunctionRunner.java | 11 +-
.../python/BaseRowPythonScalarFunctionRunner.java | 9 +-
.../python/PythonScalarFunctionRunner.java | 9 +-
.../typeutils}/BeamTypeUtils.java | 6 +-
.../typeutils}/coders/BaseRowCoder.java | 2 +-
.../typeutils}/coders/ReusableDataInputView.java | 2 +-
.../typeutils}/coders/ReusableDataOutputView.java | 2 +-
.../typeutils}/coders/RowCoder.java | 2 +-
.../org/apache/flink/python/PythonOptionsTest.java | 59 +++++
.../AbstractPythonScalarFunctionRunnerTest.java | 1 +
.../BaseRowPythonScalarFunctionRunnerTest.java | 6 +-
.../table/functions/python/BeamTypeUtilsTest.java | 5 +-
.../python/PythonScalarFunctionRunnerTest.java | 12 +-
.../BaseRowPythonScalarFunctionOperatorTest.java | 103 ++++++++
.../python/PassThroughPythonFunctionRunner.java | 75 ++++++
.../python/PythonScalarFunctionOperatorTest.java | 81 ++++++
.../PythonScalarFunctionOperatorTestBase.java | 223 +++++++++++++++++
.../typeutils}/coders/BaseRowCoderTest.java | 2 +-
.../typeutils}/coders/CoderTestBase.java | 2 +-
.../typeutils}/coders/RowCoderTest.java | 2 +-
.../util/AbstractStreamOperatorTestHarness.java | 7 +
30 files changed, 1531 insertions(+), 46 deletions(-)
diff --git a/flink-core/src/main/java/org/apache/flink/types/Row.java b/flink-core/src/main/java/org/apache/flink/types/Row.java
index b8bdbf9..aa15bf9 100644
--- a/flink-core/src/main/java/org/apache/flink/types/Row.java
+++ b/flink-core/src/main/java/org/apache/flink/types/Row.java
@@ -167,4 +167,34 @@ public class Row implements Serializable{
}
return newRow;
}
+
+ /**
+ * Creates a new Row which fields are copied from the other rows.
+ * This method does not perform a deep copy.
+ *
+ * @param first The first row being copied.
+ * @param remainings The other rows being copied.
+ * @return the joined new Row
+ */
+ public static Row join(Row first, Row... remainings) {
+ int newLength = first.fields.length;
+ for (Row remaining : remainings) {
+ newLength += remaining.fields.length;
+ }
+
+ final Row joinedRow = new Row(newLength);
+ int index = 0;
+
+ // copy the first row
+ System.arraycopy(first.fields, 0, joinedRow.fields, index, first.fields.length);
+ index += first.fields.length;
+
+ // copy the remaining rows
+ for (Row remaining : remainings) {
+ System.arraycopy(remaining.fields, 0, joinedRow.fields, index, remaining.fields.length);
+ index += remaining.fields.length;
+ }
+
+ return joinedRow;
+ }
}
diff --git a/flink-core/src/test/java/org/apache/flink/types/RowTest.java b/flink-core/src/test/java/org/apache/flink/types/RowTest.java
index 067992a..605ac9b 100644
--- a/flink-core/src/test/java/org/apache/flink/types/RowTest.java
+++ b/flink-core/src/test/java/org/apache/flink/types/RowTest.java
@@ -79,4 +79,28 @@ public class RowTest {
expected.setField(2, "hello world");
assertEquals(expected, projected);
}
+
+ @Test
+ public void testRowJoin() {
+ Row row1 = new Row(2);
+ row1.setField(0, 1);
+ row1.setField(1, "hello");
+
+ Row row2 = new Row(2);
+ row2.setField(0, null);
+ row2.setField(1, new Tuple2<>(2, "hi"));
+
+ Row row3 = new Row(1);
+ row3.setField(0, "hello world");
+
+ Row joinedRow = Row.join(row1, row2, row3);
+
+ Row expected = new Row(5);
+ expected.setField(0, 1);
+ expected.setField(1, "hello");
+ expected.setField(2, null);
+ expected.setField(3, new Tuple2<>(2, "hi"));
+ expected.setField(4, "hello world");
+ assertEquals(expected, joinedRow);
+ }
}
diff --git a/flink-python/pom.xml b/flink-python/pom.xml
index fcbdb0c..ed16daa 100644
--- a/flink-python/pom.xml
+++ b/flink-python/pom.xml
@@ -74,6 +74,12 @@ under the License.
<version>${project.version}</version>
<scope>provided</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-table-planner-blink_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <scope>provided</scope>
+ </dependency>
<!-- Beam dependencies -->
diff --git a/flink-python/src/main/java/org/apache/flink/python/AbstractPythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/python/AbstractPythonFunctionRunner.java
index 498e090..73c7535 100644
--- a/flink-python/src/main/java/org/apache/flink/python/AbstractPythonFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/python/AbstractPythonFunctionRunner.java
@@ -49,6 +49,7 @@ import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
+import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;
@@ -92,9 +93,9 @@ public abstract class AbstractPythonFunctionRunner<IN, OUT> implements PythonFun
private final StateRequestHandler stateRequestHandler;
/**
- * Temporary directory to store the retrieval token.
+ * Temporary directories to store the retrieval token.
*/
- private final String tempDir;
+ private final String[] tempDirs;
/**
* The file of the retrieval token representing the entirety of the staged artifacts.
@@ -146,12 +147,12 @@ public abstract class AbstractPythonFunctionRunner<IN, OUT> implements PythonFun
FnDataReceiver<OUT> resultReceiver,
PythonEnv pythonEnv,
StateRequestHandler stateRequestHandler,
- String tempDir) {
+ String[] tempDirs) {
this.taskName = Preconditions.checkNotNull(taskName);
this.resultReceiver = Preconditions.checkNotNull(resultReceiver);
this.pythonEnv = Preconditions.checkNotNull(pythonEnv);
this.stateRequestHandler = Preconditions.checkNotNull(stateRequestHandler);
- this.tempDir = Preconditions.checkNotNull(tempDir);
+ this.tempDirs = Preconditions.checkNotNull(tempDirs);
}
@Override
@@ -249,23 +250,22 @@ public abstract class AbstractPythonFunctionRunner<IN, OUT> implements PythonFun
}
private String createEmptyRetrievalToken() throws Exception {
- // try to find a unique file name for the spilling channel
+ // try to find a unique file name for the retrieval token
final Random rnd = new Random();
for (int attempt = 0; attempt < 10; attempt++) {
- retrievalToken = new File(tempDir, randomString(rnd) + ".json");
- if (!retrievalToken.exists()) {
- break;
+ String directory = tempDirs[rnd.nextInt(tempDirs.length)];
+ retrievalToken = new File(directory, randomString(rnd) + ".json");
+ if (retrievalToken.createNewFile()) {
+ final DataOutputStream dos = new DataOutputStream(new FileOutputStream(retrievalToken));
+ dos.writeBytes("{\"manifest\": {}}");
+ dos.flush();
+ dos.close();
+ return retrievalToken.getAbsolutePath();
}
}
- if (retrievalToken.exists()) {
- throw new IOException(
- "Could not find a unique file name in '" + tempDir + "' for retrieval token.");
- }
- final DataOutputStream dos = new DataOutputStream(new FileOutputStream(retrievalToken));
- dos.writeBytes("{\"manifest\": {}}");
- dos.flush();
- dos.close();
- return retrievalToken.getAbsolutePath();
+
+ throw new IOException(
+ "Could not find a unique file name in '" + Arrays.toString(tempDirs) + "' for retrieval token.");
}
private static String randomString(Random random) {
diff --git a/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java b/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java
new file mode 100644
index 0000000..6c0ee52
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/python/PythonOptions.java
@@ -0,0 +1,48 @@
+/*
+ * 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.python;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
+
+/**
+ * Configuration options for the Python API.
+ */
+@PublicEvolving
+public class PythonOptions {
+
+ /**
+ * The maximum number of elements to include in a bundle.
+ */
+ public static final ConfigOption<Integer> MAX_BUNDLE_SIZE = ConfigOptions
+ .key("python.fn-execution.bundle.size")
+ .defaultValue(1000)
+ .withDescription("The maximum number of elements to include in a bundle. The elements " +
+ "are processed asynchronously. One bundle of elements are processed before " +
+ "processing the next bundle of elements");
+
+ /**
+ * The maximum time to wait before finalising a bundle (in milliseconds).
+ */
+ public static final ConfigOption<Long> MAX_BUNDLE_TIME_MILLS = ConfigOptions
+ .key("python.fn-execution.bundle.time")
+ .defaultValue(1000L)
+ .withDescription("The maximum time to wait before finishing a bundle (in milliseconds).");
+}
diff --git a/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java
new file mode 100644
index 0000000..1c869db
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/streaming/api/operators/python/AbstractPythonFunctionOperator.java
@@ -0,0 +1,278 @@
+/*
+ * 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.streaming.api.operators.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
+import org.apache.flink.streaming.api.operators.BoundedOneInput;
+import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+
+import java.util.Map;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Base class for all stream operators to execute Python functions.
+ */
+@Internal
+public abstract class AbstractPythonFunctionOperator<IN, OUT>
+ extends AbstractStreamOperator<OUT>
+ implements OneInputStreamOperator<IN, OUT>, BoundedOneInput {
+
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * The {@link PythonFunctionRunner} which is responsible for Python user-defined function execution.
+ */
+ private transient PythonFunctionRunner<IN> pythonFunctionRunner;
+
+ /**
+ * Use an AtomicBoolean because we start/stop bundles by a timer thread.
+ */
+ private transient AtomicBoolean bundleStarted;
+
+ /**
+ * Number of processed elements in the current bundle.
+ */
+ private transient int elementCount;
+
+ /**
+ * Max number of elements to include in a bundle.
+ */
+ private transient int maxBundleSize;
+
+ /**
+ * Max duration of a bundle.
+ */
+ private transient long maxBundleTimeMills;
+
+ /**
+ * Time that the last bundle was finished.
+ */
+ private transient long lastFinishBundleTime;
+
+ /**
+ * A timer that finishes the current bundle after a fixed amount of time.
+ */
+ private transient ScheduledFuture<?> checkFinishBundleTimer;
+
+ /**
+ * Callback to be executed after the current bundle was finished.
+ */
+ private transient Runnable bundleFinishedCallback;
+
+ @Override
+ public void open() throws Exception {
+ try {
+ this.bundleStarted = new AtomicBoolean(false);
+
+ Map<String, String> jobParams = getExecutionConfig().getGlobalJobParameters().toMap();
+
+ this.maxBundleSize = Integer.valueOf(jobParams.getOrDefault(
+ PythonOptions.MAX_BUNDLE_SIZE.key(),
+ String.valueOf(PythonOptions.MAX_BUNDLE_SIZE.defaultValue())));
+ if (this.maxBundleSize <= 0) {
+ this.maxBundleSize = PythonOptions.MAX_BUNDLE_SIZE.defaultValue();
+ LOG.error("Invalid value for the maximum bundle size. Using default value of " +
+ this.maxBundleSize + '.');
+ } else {
+ LOG.info("The maximum bundle size is configured to {}.", this.maxBundleSize);
+ }
+
+ this.maxBundleTimeMills = Long.valueOf(jobParams.getOrDefault(
+ PythonOptions.MAX_BUNDLE_TIME_MILLS.key(),
+ String.valueOf(PythonOptions.MAX_BUNDLE_TIME_MILLS.defaultValue())));
+ if (this.maxBundleTimeMills <= 0L) {
+ this.maxBundleTimeMills = PythonOptions.MAX_BUNDLE_TIME_MILLS.defaultValue();
+ LOG.error("Invalid value for the maximum bundle time. Using default value of " +
+ this.maxBundleTimeMills + '.');
+ } else {
+ LOG.info("The maximum bundle time is configured to {} milliseconds.", this.maxBundleTimeMills);
+ }
+
+ this.pythonFunctionRunner = createPythonFunctionRunner();
+ this.pythonFunctionRunner.open();
+
+ this.elementCount = 0;
+ this.lastFinishBundleTime = getProcessingTimeService().getCurrentProcessingTime();
+
+ // Schedule timer to check timeout of finish bundle.
+ long bundleCheckPeriod = Math.max(this.maxBundleTimeMills, 1);
+ this.checkFinishBundleTimer =
+ getProcessingTimeService()
+ .scheduleAtFixedRate(
+ // ProcessingTimeService callbacks are executed under the checkpointing lock
+ timestamp -> checkInvokeFinishBundleByTime(), bundleCheckPeriod, bundleCheckPeriod);
+ } finally {
+ super.open();
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ try {
+ invokeFinishBundle();
+ } finally {
+ super.close();
+ }
+ }
+
+ @Override
+ public void dispose() throws Exception {
+ try {
+ if (checkFinishBundleTimer != null) {
+ checkFinishBundleTimer.cancel(true);
+ checkFinishBundleTimer = null;
+ }
+ if (pythonFunctionRunner != null) {
+ pythonFunctionRunner.close();
+ pythonFunctionRunner = null;
+ }
+ } finally {
+ super.dispose();
+ }
+ }
+
+ @Override
+ public void endInput() throws Exception {
+ invokeFinishBundle();
+ }
+
+ @Override
+ public void processElement(StreamRecord<IN> element) throws Exception {
+ checkInvokeStartBundle();
+ pythonFunctionRunner.processElement(element.getValue());
+ checkInvokeFinishBundleByCount();
+ }
+
+ @Override
+ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
+ try {
+ // Ensures that no new bundle gets started
+ invokeFinishBundle();
+ } finally {
+ super.prepareSnapshotPreBarrier(checkpointId);
+ }
+ }
+
+ @Override
+ public void processWatermark(Watermark mark) throws Exception {
+ // Due to the asynchronous communication with the SDK harness,
+ // a bundle might still be in progress and not all items have
+ // yet been received from the SDK harness. If we just set this
+ // watermark as the new output watermark, we could violate the
+ // order of the records, i.e. pending items in the SDK harness
+ // could become "late" although they were "on time".
+ //
+ // We can solve this problem using one of the following options:
+ //
+ // 1) Finish the current bundle and emit this watermark as the
+ // new output watermark. Finishing the bundle ensures that
+ // all the items have been processed by the SDK harness and
+ // the execution results sent to the downstream operator.
+ //
+ // 2) Hold on the output watermark for as long as the current
+ // bundle has not been finished. We have to remember to manually
+ // finish the bundle in case we receive the final watermark.
+ // To avoid latency, we should process this watermark again as
+ // soon as the current bundle is finished.
+ //
+ // Approach 1) is the easiest and gives better latency, yet 2)
+ // gives better throughput due to the bundle not getting cut on
+ // every watermark. So we have implemented 2) below.
+ if (mark.getTimestamp() == Long.MAX_VALUE) {
+ invokeFinishBundle();
+ super.processWatermark(mark);
+ } else if (!bundleStarted.get()) {
+ // forward the watermark immediately if the bundle is already finished.
+ super.processWatermark(mark);
+ } else {
+ // It is not safe to advance the output watermark yet, so add a hold on the current
+ // output watermark.
+ bundleFinishedCallback =
+ () -> {
+ try {
+ // at this point the bundle is finished, allow the watermark to pass
+ super.processWatermark(mark);
+ } catch (Exception e) {
+ throw new RuntimeException(
+ "Failed to process watermark after finished bundle.", e);
+ }
+ };
+ }
+ }
+
+ /**
+ * Creates the {@link PythonFunctionRunner} which is responsible for Python user-defined function execution.
+ */
+ public abstract PythonFunctionRunner<IN> createPythonFunctionRunner();
+
+ /**
+ * Sends the execution results to the downstream operator.
+ */
+ public abstract void emitResults();
+
+ /**
+ * Checks whether to invoke startBundle.
+ */
+ private void checkInvokeStartBundle() throws Exception {
+ if (bundleStarted.compareAndSet(false, true)) {
+ pythonFunctionRunner.startBundle();
+ }
+ }
+
+ /**
+ * Checks whether to invoke finishBundle by elements count. Called in processElement.
+ */
+ private void checkInvokeFinishBundleByCount() throws Exception {
+ elementCount++;
+ if (elementCount >= maxBundleSize) {
+ invokeFinishBundle();
+ }
+ }
+
+ /**
+ * Checks whether to invoke finishBundle by timeout.
+ */
+ private void checkInvokeFinishBundleByTime() throws Exception {
+ long now = getProcessingTimeService().getCurrentProcessingTime();
+ if (now - lastFinishBundleTime >= maxBundleTimeMills) {
+ invokeFinishBundle();
+ }
+ }
+
+ private void invokeFinishBundle() throws Exception {
+ if (bundleStarted.compareAndSet(true, false)) {
+ pythonFunctionRunner.finishBundle();
+
+ emitResults();
+ elementCount = 0;
+ lastFinishBundleTime = getProcessingTimeService().getCurrentProcessingTime();
+ // callback only after current bundle was fully finalized
+ if (bundleFinishedCallback != null) {
+ bundleFinishedCallback.run();
+ bundleFinishedCallback = null;
+ }
+ }
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractPythonScalarFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractPythonScalarFunctionOperator.java
new file mode 100644
index 0000000..4ce8a9c
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/AbstractPythonScalarFunctionOperator.java
@@ -0,0 +1,200 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.streaming.api.operators.python.AbstractPythonFunctionOperator;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+import java.util.Arrays;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.stream.Collectors;
+
+/**
+ * Base class for all stream operators to execute Python {@link ScalarFunction}s. It executes the Python
+ * {@link ScalarFunction}s in separate Python execution environment.
+ *
+ * <p>The inputs are assumed as the following format:
+ * {{{
+ * +------------------+--------------+
+ * | forwarded fields | extra fields |
+ * +------------------+--------------+
+ * }}}.
+ *
+ * <p>The Python UDFs may take input columns directly from the input row or the execution result of Java UDFs:
+ * 1) The input columns from the input row can be referred from the 'forwarded fields';
+ * 2) The Java UDFs will be computed and the execution results can be referred from the 'extra fields'.
+ *
+ * <p>The outputs will be as the following format:
+ * {{{
+ * +------------------+-------------------------+
+ * | forwarded fields | scalar function results |
+ * +------------------+-------------------------+
+ * }}}.
+ *
+ * @param <IN> Type of the input elements.
+ * @param <OUT> Type of the output elements.
+ * @param <UDFIN> Type of the UDF input type.
+ * @param <UDFOUT> Type of the UDF input type.
+ */
+@Internal
+public abstract class AbstractPythonScalarFunctionOperator<IN, OUT, UDFIN, UDFOUT>
+ extends AbstractPythonFunctionOperator<IN, OUT> {
+
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * The Python {@link ScalarFunction}s to be executed.
+ */
+ protected final PythonFunctionInfo[] scalarFunctions;
+
+ /**
+ * The input logical type.
+ */
+ protected final RowType inputType;
+
+ /**
+ * The output logical type.
+ */
+ protected final RowType outputType;
+
+ /**
+ * The offsets of udf inputs.
+ */
+ protected final int[] udfInputOffsets;
+
+ /**
+ * The number of forwarded fields in the input element.
+ */
+ protected final int forwardedFieldCnt;
+
+ /**
+ * The udf input logical type.
+ */
+ protected transient RowType udfInputType;
+
+ /**
+ * The udf output logical type.
+ */
+ protected transient RowType udfOutputType;
+
+ /**
+ * The queue holding the input elements for which the execution results have not been received.
+ */
+ protected transient LinkedBlockingQueue<IN> forwardedInputQueue;
+
+ /**
+ * The queue holding the user-defined function execution results. The execution results are in
+ * the same order as the input elements.
+ */
+ protected transient LinkedBlockingQueue<UDFOUT> udfResultQueue;
+
+ AbstractPythonScalarFunctionOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ this.scalarFunctions = Preconditions.checkNotNull(scalarFunctions);
+ this.inputType = Preconditions.checkNotNull(inputType);
+ this.outputType = Preconditions.checkNotNull(outputType);
+ this.udfInputOffsets = Preconditions.checkNotNull(udfInputOffsets);
+ this.forwardedFieldCnt = forwardedFieldCnt;
+ }
+
+ @Override
+ public void open() throws Exception {
+ forwardedInputQueue = new LinkedBlockingQueue<>();
+ udfResultQueue = new LinkedBlockingQueue<>();
+ udfInputType = new RowType(
+ Arrays.stream(udfInputOffsets)
+ .mapToObj(i -> inputType.getFields().get(i))
+ .collect(Collectors.toList()));
+ udfOutputType = new RowType(outputType.getFields().subList(forwardedFieldCnt, outputType.getFieldCount()));
+ super.open();
+ }
+
+ @Override
+ public void processElement(StreamRecord<IN> element) throws Exception {
+ bufferInput(element.getValue());
+ super.processElement(element);
+ emitResults();
+ }
+
+ @Override
+ public PythonFunctionRunner<IN> createPythonFunctionRunner() {
+ final FnDataReceiver<UDFOUT> udfResultReceiver = input -> {
+ // handover to queue, do not block the result receiver thread
+ udfResultQueue.put(input);
+ };
+
+ return new ProjectUdfInputPythonScalarFunctionRunner(createPythonFunctionRunner(udfResultReceiver));
+ }
+
+ /**
+ * Buffers the specified input, it will be used to construct
+ * the operator result together with the udf execution result.
+ */
+ public abstract void bufferInput(IN input);
+
+ public abstract UDFIN getUdfInput(IN element);
+
+ public abstract PythonFunctionRunner<UDFIN> createPythonFunctionRunner(FnDataReceiver<UDFOUT> resultReceiver);
+
+ private class ProjectUdfInputPythonScalarFunctionRunner implements PythonFunctionRunner<IN> {
+
+ private final PythonFunctionRunner<UDFIN> pythonFunctionRunner;
+
+ ProjectUdfInputPythonScalarFunctionRunner(PythonFunctionRunner<UDFIN> pythonFunctionRunner) {
+ this.pythonFunctionRunner = pythonFunctionRunner;
+ }
+
+ @Override
+ public void open() throws Exception {
+ pythonFunctionRunner.open();
+ }
+
+ @Override
+ public void close() throws Exception {
+ pythonFunctionRunner.close();
+ }
+
+ @Override
+ public void startBundle() throws Exception {
+ pythonFunctionRunner.startBundle();
+ }
+
+ @Override
+ public void finishBundle() throws Exception {
+ pythonFunctionRunner.finishBundle();
+ }
+
+ @Override
+ public void processElement(IN element) throws Exception {
+ pythonFunctionRunner.processElement(getUdfInput(element));
+ }
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperator.java
new file mode 100644
index 0000000..2f53808
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperator.java
@@ -0,0 +1,178 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.table.api.TableConfig;
+import org.apache.flink.table.dataformat.BaseRow;
+import org.apache.flink.table.dataformat.BinaryRow;
+import org.apache.flink.table.dataformat.JoinedRow;
+import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.planner.codegen.CodeGeneratorContext;
+import org.apache.flink.table.planner.codegen.ProjectionCodeGenerator;
+import org.apache.flink.table.runtime.generated.GeneratedProjection;
+import org.apache.flink.table.runtime.generated.Projection;
+import org.apache.flink.table.runtime.runners.python.BaseRowPythonScalarFunctionRunner;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.util.Collector;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+/**
+ * The Python {@link ScalarFunction} operator for the blink planner.
+ */
+@Internal
+public class BaseRowPythonScalarFunctionOperator
+ extends AbstractPythonScalarFunctionOperator<BaseRow, BaseRow, BaseRow, BaseRow> {
+
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * The collector used to collect records.
+ */
+ private transient StreamRecordBaseRowWrappingCollector baseRowWrapper;
+
+ /**
+ * The JoinedRow reused holding the execution result.
+ */
+ private transient JoinedRow reuseJoinedRow;
+
+ /**
+ * The Projection which projects the forwarded fields from the input row.
+ */
+ private transient Projection<BaseRow, BinaryRow> forwardedFieldProjection;
+
+ /**
+ * The Projection which projects the udf input fields from the input row.
+ */
+ private transient Projection<BaseRow, BinaryRow> udfInputProjection;
+
+ public BaseRowPythonScalarFunctionOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ super(scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFieldCnt);
+ }
+
+ @Override
+ public void open() throws Exception {
+ super.open();
+ baseRowWrapper = new StreamRecordBaseRowWrappingCollector(output);
+ reuseJoinedRow = new JoinedRow();
+
+ udfInputProjection = createUdfInputProjection();
+ forwardedFieldProjection = createForwardedFieldProjection();
+ }
+
+ @Override
+ public void bufferInput(BaseRow input) {
+ // always copy the projection result as the generated Projection reuses the projection result
+ BaseRow forwardedFields = forwardedFieldProjection.apply(input).copy();
+ forwardedFields.setHeader(input.getHeader());
+ forwardedInputQueue.add(forwardedFields);
+ }
+
+ @Override
+ public BaseRow getUdfInput(BaseRow element) {
+ return udfInputProjection.apply(element);
+ }
+
+ @Override
+ @SuppressWarnings("ConstantConditions")
+ public void emitResults() {
+ BaseRow udfResult;
+ while ((udfResult = udfResultQueue.poll()) != null) {
+ BaseRow input = forwardedInputQueue.poll();
+ reuseJoinedRow.setHeader(input.getHeader());
+ baseRowWrapper.collect(reuseJoinedRow.replace(input, udfResult));
+ }
+ }
+
+ @Override
+ public PythonFunctionRunner<BaseRow> createPythonFunctionRunner(FnDataReceiver<BaseRow> resultReceiver) {
+ return new BaseRowPythonScalarFunctionRunner(
+ getRuntimeContext().getTaskName(),
+ resultReceiver,
+ scalarFunctions,
+ scalarFunctions[0].getPythonFunction().getPythonEnv(),
+ udfInputType,
+ udfOutputType,
+ getContainingTask().getEnvironment().getTaskManagerInfo().getTmpDirectories());
+ }
+
+ private Projection<BaseRow, BinaryRow> createUdfInputProjection() {
+ final GeneratedProjection generatedProjection = ProjectionCodeGenerator.generateProjection(
+ CodeGeneratorContext.apply(new TableConfig()),
+ "UdfInputProjection",
+ inputType,
+ udfInputType,
+ udfInputOffsets);
+ // noinspection unchecked
+ return generatedProjection.newInstance(Thread.currentThread().getContextClassLoader());
+ }
+
+ private Projection<BaseRow, BinaryRow> createForwardedFieldProjection() {
+ final int[] fields = new int[forwardedFieldCnt];
+ for (int i = 0; i < fields.length; i++) {
+ fields[i] = i;
+ }
+
+ final RowType forwardedFieldType = new RowType(inputType.getFields().subList(0, forwardedFieldCnt));
+ final GeneratedProjection generatedProjection = ProjectionCodeGenerator.generateProjection(
+ CodeGeneratorContext.apply(new TableConfig()),
+ "ForwardedFieldProjection",
+ inputType,
+ forwardedFieldType,
+ fields);
+ // noinspection unchecked
+ return generatedProjection.newInstance(Thread.currentThread().getContextClassLoader());
+ }
+
+ /**
+ * The collector is used to convert a {@link BaseRow} to a {@link StreamRecord}.
+ */
+ private static class StreamRecordBaseRowWrappingCollector implements Collector<BaseRow> {
+
+ private final Collector<StreamRecord<BaseRow>> out;
+
+ /**
+ * For Table API & SQL jobs, the timestamp field is not used.
+ */
+ private final StreamRecord<BaseRow> reuseStreamRecord = new StreamRecord<>(null);
+
+ StreamRecordBaseRowWrappingCollector(Collector<StreamRecord<BaseRow>> out) {
+ this.out = out;
+ }
+
+ @Override
+ public void collect(BaseRow record) {
+ out.collect(reuseStreamRecord.replace(record));
+ }
+
+ @Override
+ public void close() {
+ out.close();
+ }
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperator.java b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperator.java
new file mode 100644
index 0000000..9961b16
--- /dev/null
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperator.java
@@ -0,0 +1,158 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.runners.python.PythonScalarFunctionRunner;
+import org.apache.flink.table.runtime.types.CRow;
+import org.apache.flink.table.runtime.types.CRowTypeInfo;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.table.types.utils.TypeConversions;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.Collector;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+/**
+ * The Python {@link ScalarFunction} operator for the legacy planner.
+ */
+@Internal
+public class PythonScalarFunctionOperator extends AbstractPythonScalarFunctionOperator<CRow, CRow, Row, Row> {
+
+ private static final long serialVersionUID = 1L;
+
+ /**
+ * The collector used to collect records.
+ */
+ private transient StreamRecordCRowWrappingCollector cRowWrapper;
+
+ /**
+ * The type serializer for the forwarded fields.
+ */
+ private transient TypeSerializer<CRow> forwardedInputSerializer;
+
+ public PythonScalarFunctionOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ super(scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFieldCnt);
+ }
+
+ @Override
+ public void open() throws Exception {
+ super.open();
+ this.cRowWrapper = new StreamRecordCRowWrappingCollector(output);
+
+ CRowTypeInfo forwardedInputTypeInfo = new CRowTypeInfo(new RowTypeInfo(
+ inputType.getFields().stream()
+ .limit(forwardedFieldCnt)
+ .map(RowType.RowField::getType)
+ .map(TypeConversions::fromLogicalToDataType)
+ .map(TypeConversions::fromDataTypeToLegacyInfo)
+ .toArray(TypeInformation[]::new)));
+ forwardedInputSerializer = forwardedInputTypeInfo.createSerializer(getExecutionConfig());
+ }
+
+ @Override
+ public void bufferInput(CRow input) {
+ CRow forwardedFieldsRow = new CRow(getForwardedRow(input.row()), input.change());
+ if (getExecutionConfig().isObjectReuseEnabled()) {
+ forwardedFieldsRow = forwardedInputSerializer.copy(forwardedFieldsRow);
+ }
+ forwardedInputQueue.add(forwardedFieldsRow);
+ }
+
+ @Override
+ public Row getUdfInput(CRow element) {
+ return Row.project(element.row(), udfInputOffsets);
+ }
+
+ @Override
+ @SuppressWarnings("ConstantConditions")
+ public void emitResults() {
+ Row udfResult;
+ while ((udfResult = udfResultQueue.poll()) != null) {
+ CRow input = forwardedInputQueue.poll();
+ cRowWrapper.setChange(input.change());
+ cRowWrapper.collect(Row.join(input.row(), udfResult));
+ }
+ }
+
+ @Override
+ public PythonFunctionRunner<Row> createPythonFunctionRunner(FnDataReceiver<Row> resultReceiver) {
+ return new PythonScalarFunctionRunner(
+ getRuntimeContext().getTaskName(),
+ resultReceiver,
+ scalarFunctions,
+ scalarFunctions[0].getPythonFunction().getPythonEnv(),
+ udfInputType,
+ udfOutputType,
+ getContainingTask().getEnvironment().getTaskManagerInfo().getTmpDirectories());
+ }
+
+ private Row getForwardedRow(Row input) {
+ Row row = new Row(forwardedFieldCnt);
+ for (int i = 0; i < row.getArity(); i++) {
+ row.setField(i, input.getField(i));
+ }
+ return row;
+ }
+
+ /**
+ * The collector is used to convert a {@link Row} to a {@link CRow}.
+ */
+ private static class StreamRecordCRowWrappingCollector implements Collector<Row> {
+
+ private final Collector<StreamRecord<CRow>> out;
+ private final CRow reuseCRow = new CRow();
+
+ /**
+ * For Table API & SQL jobs, the timestamp field is not used.
+ */
+ private final StreamRecord<CRow> reuseStreamRecord = new StreamRecord<>(reuseCRow);
+
+ StreamRecordCRowWrappingCollector(Collector<StreamRecord<CRow>> out) {
+ this.out = out;
+ }
+
+ public void setChange(boolean change) {
+ this.reuseCRow.change_$eq(change);
+ }
+
+ @Override
+ public void collect(Row record) {
+ reuseCRow.row_$eq(record);
+ out.collect(reuseStreamRecord);
+ }
+
+ @Override
+ public void close() {
+ out.close();
+ }
+ }
+}
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/AbstractPythonScalarFunctionRunner.java
similarity index 95%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunner.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/AbstractPythonScalarFunctionRunner.java
index 104b058..2679156 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/AbstractPythonScalarFunctionRunner.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python;
+package org.apache.flink.table.runtime.runners.python;
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
@@ -24,6 +24,9 @@ import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.AbstractPythonFunctionRunner;
import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.typeutils.BeamTypeUtils;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.util.Preconditions;
@@ -78,8 +81,8 @@ public abstract class AbstractPythonScalarFunctionRunner<IN, OUT> extends Abstra
PythonEnv pythonEnv,
RowType inputType,
RowType outputType,
- String tempDir) {
- super(taskName, resultReceiver, pythonEnv, StateRequestHandler.unsupported(), tempDir);
+ String[] tempDirs) {
+ super(taskName, resultReceiver, pythonEnv, StateRequestHandler.unsupported(), tempDirs);
this.scalarFunctions = Preconditions.checkNotNull(scalarFunctions);
this.inputType = Preconditions.checkNotNull(inputType);
this.outputType = Preconditions.checkNotNull(outputType);
@@ -164,7 +167,7 @@ public abstract class AbstractPythonScalarFunctionRunner<IN, OUT> extends Abstra
* Gets the proto representation of the Python user-defined functions to be executed.
*/
@VisibleForTesting
- FlinkFnApi.UserDefinedFunctions getUserDefinedFunctionsProto() {
+ public FlinkFnApi.UserDefinedFunctions getUserDefinedFunctionsProto() {
FlinkFnApi.UserDefinedFunctions.Builder builder = FlinkFnApi.UserDefinedFunctions.newBuilder();
for (PythonFunctionInfo pythonFunctionInfo : scalarFunctions) {
builder.addUdfs(getUserDefinedFunctionProto(pythonFunctionInfo));
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/BaseRowPythonScalarFunctionRunner.java
similarity index 87%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunner.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/BaseRowPythonScalarFunctionRunner.java
index 5f4cbce..4bb5fc8 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/BaseRowPythonScalarFunctionRunner.java
@@ -16,12 +16,15 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python;
+package org.apache.flink.table.runtime.runners.python;
import org.apache.flink.annotation.Internal;
import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.table.dataformat.BaseRow;
import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.typeutils.BeamTypeUtils;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.sdk.coders.Coder;
@@ -41,8 +44,8 @@ public class BaseRowPythonScalarFunctionRunner extends AbstractPythonScalarFunct
PythonEnv pythonEnv,
RowType inputType,
RowType outputType,
- String tempDir) {
- super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDir);
+ String[] tempDirs) {
+ super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDirs);
}
@Override
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/PythonScalarFunctionRunner.java
similarity index 86%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunner.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/PythonScalarFunctionRunner.java
index 92f2dd0..8edba3a 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunner.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/runners/python/PythonScalarFunctionRunner.java
@@ -16,11 +16,14 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python;
+package org.apache.flink.table.runtime.runners.python;
import org.apache.flink.annotation.Internal;
import org.apache.flink.python.PythonFunctionRunner;
import org.apache.flink.table.functions.ScalarFunction;
+import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.typeutils.BeamTypeUtils;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.types.Row;
@@ -41,8 +44,8 @@ public class PythonScalarFunctionRunner extends AbstractPythonScalarFunctionRunn
PythonEnv pythonEnv,
RowType inputType,
RowType outputType,
- String tempDir) {
- super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDir);
+ String[] tempDirs) {
+ super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDirs);
}
@Override
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/BeamTypeUtils.java b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/BeamTypeUtils.java
similarity index 98%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/BeamTypeUtils.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/BeamTypeUtils.java
index ac3e989..710529c 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/BeamTypeUtils.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/BeamTypeUtils.java
@@ -16,12 +16,12 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python;
+package org.apache.flink.table.runtime.typeutils;
import org.apache.flink.annotation.Internal;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.table.functions.python.coders.BaseRowCoder;
-import org.apache.flink.table.functions.python.coders.RowCoder;
+import org.apache.flink.table.runtime.typeutils.coders.BaseRowCoder;
+import org.apache.flink.table.runtime.typeutils.coders.RowCoder;
import org.apache.flink.table.types.logical.AnyType;
import org.apache.flink.table.types.logical.ArrayType;
import org.apache.flink.table.types.logical.BigIntType;
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/BaseRowCoder.java b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoder.java
similarity index 98%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/coders/BaseRowCoder.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoder.java
index 8a07c22..9834c50 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/BaseRowCoder.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoder.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.annotation.Internal;
import org.apache.flink.core.memory.DataOutputView;
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataInputView.java b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataInputView.java
similarity index 95%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataInputView.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataInputView.java
index 5beaa69..d71cf3b 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataInputView.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataInputView.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.annotation.Internal;
import org.apache.flink.core.memory.DataInputView;
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataOutputView.java b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataOutputView.java
similarity index 95%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataOutputView.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataOutputView.java
index 9aabb4f..6128044 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/ReusableDataOutputView.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/ReusableDataOutputView.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.annotation.Internal;
import org.apache.flink.core.memory.DataOutputView;
diff --git a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/RowCoder.java b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/RowCoder.java
similarity index 98%
rename from flink-python/src/main/java/org/apache/flink/table/functions/python/coders/RowCoder.java
rename to flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/RowCoder.java
index 5b1e8e0..b72d5f1 100644
--- a/flink-python/src/main/java/org/apache/flink/table/functions/python/coders/RowCoder.java
+++ b/flink-python/src/main/java/org/apache/flink/table/runtime/typeutils/coders/RowCoder.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.annotation.Internal;
import org.apache.flink.types.Row;
diff --git a/flink-python/src/test/java/org/apache/flink/python/PythonOptionsTest.java b/flink-python/src/test/java/org/apache/flink/python/PythonOptionsTest.java
new file mode 100644
index 0000000..74072dc
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/python/PythonOptionsTest.java
@@ -0,0 +1,59 @@
+/*
+ * 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.python;
+
+import org.apache.flink.configuration.Configuration;
+
+import org.junit.Test;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertThat;
+
+/**
+ * Test all configurations can be set using configuration.
+ */
+public class PythonOptionsTest {
+
+ @Test
+ public void testBundleSize() {
+ final Configuration configuration = new Configuration();
+ final int defaultBundleSize = configuration.getInteger(PythonOptions.MAX_BUNDLE_SIZE);
+ assertThat(defaultBundleSize, is(equalTo(PythonOptions.MAX_BUNDLE_SIZE.defaultValue())));
+
+ final int expectedBundleSize = 100;
+ configuration.setInteger(PythonOptions.MAX_BUNDLE_SIZE, expectedBundleSize);
+
+ final int actualBundleSize = configuration.getInteger(PythonOptions.MAX_BUNDLE_SIZE);
+ assertThat(actualBundleSize, is(equalTo(expectedBundleSize)));
+ }
+
+ @Test
+ public void testBundleTime() {
+ final Configuration configuration = new Configuration();
+ final long defaultBundleTime = configuration.getLong(PythonOptions.MAX_BUNDLE_TIME_MILLS);
+ assertThat(defaultBundleTime, is(equalTo(PythonOptions.MAX_BUNDLE_TIME_MILLS.defaultValue())));
+
+ final long expectedBundleTime = 100;
+ configuration.setLong(PythonOptions.MAX_BUNDLE_TIME_MILLS, expectedBundleTime);
+
+ final long actualBundleSize = configuration.getLong(PythonOptions.MAX_BUNDLE_TIME_MILLS);
+ assertThat(actualBundleSize, is(equalTo(expectedBundleTime)));
+ }
+}
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunnerTest.java b/flink-python/src/test/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunnerTest.java
index b811d78..c53f69a 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunnerTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/functions/python/AbstractPythonScalarFunctionRunnerTest.java
@@ -18,6 +18,7 @@
package org.apache.flink.table.functions.python;
+import org.apache.flink.table.runtime.runners.python.AbstractPythonScalarFunctionRunner;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.RowType;
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunnerTest.java b/flink-python/src/test/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunnerTest.java
index 4e8e655..02f5b74 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunnerTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/functions/python/BaseRowPythonScalarFunctionRunnerTest.java
@@ -19,7 +19,9 @@
package org.apache.flink.table.functions.python;
import org.apache.flink.table.dataformat.BaseRow;
-import org.apache.flink.table.functions.python.coders.BaseRowCoder;
+import org.apache.flink.table.runtime.runners.python.AbstractPythonScalarFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.BaseRowPythonScalarFunctionRunner;
+import org.apache.flink.table.runtime.typeutils.coders.BaseRowCoder;
import org.apache.flink.table.types.logical.RowType;
import org.apache.beam.sdk.coders.Coder;
@@ -123,6 +125,6 @@ public class BaseRowPythonScalarFunctionRunnerTest extends AbstractPythonScalarF
pythonEnv,
inputType,
outputType,
- System.getProperty("java.io.tmpdir"));
+ new String[] {System.getProperty("java.io.tmpdir")});
}
}
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/BeamTypeUtilsTest.java b/flink-python/src/test/java/org/apache/flink/table/functions/python/BeamTypeUtilsTest.java
index 7acf7b3..0f3f308 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/BeamTypeUtilsTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/functions/python/BeamTypeUtilsTest.java
@@ -19,8 +19,9 @@
package org.apache.flink.table.functions.python;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.table.functions.python.coders.BaseRowCoder;
-import org.apache.flink.table.functions.python.coders.RowCoder;
+import org.apache.flink.table.runtime.typeutils.BeamTypeUtils;
+import org.apache.flink.table.runtime.typeutils.coders.BaseRowCoder;
+import org.apache.flink.table.runtime.typeutils.coders.RowCoder;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.RowType;
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunnerTest.java b/flink-python/src/test/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunnerTest.java
index c3b63bb..6a9398a 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunnerTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/functions/python/PythonScalarFunctionRunnerTest.java
@@ -19,7 +19,9 @@
package org.apache.flink.table.functions.python;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
-import org.apache.flink.table.functions.python.coders.RowCoder;
+import org.apache.flink.table.runtime.runners.python.AbstractPythonScalarFunctionRunner;
+import org.apache.flink.table.runtime.runners.python.PythonScalarFunctionRunner;
+import org.apache.flink.table.runtime.typeutils.coders.RowCoder;
import org.apache.flink.table.types.logical.BigIntType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.types.Row;
@@ -251,7 +253,7 @@ public class PythonScalarFunctionRunnerTest extends AbstractPythonScalarFunction
pythonEnv,
inputType,
outputType,
- System.getProperty("java.io.tmpdir"));
+ new String[] {System.getProperty("java.io.tmpdir")});
}
private AbstractPythonScalarFunctionRunner<Row, Row> createUDFRunner(
@@ -274,7 +276,7 @@ public class PythonScalarFunctionRunnerTest extends AbstractPythonScalarFunction
rowType,
rowType,
jobBundleFactory,
- System.getProperty("java.io.tmpdir"));
+ new String[] {System.getProperty("java.io.tmpdir")});
}
private static class PythonScalarFunctionRunnerTestHarness extends PythonScalarFunctionRunner {
@@ -288,8 +290,8 @@ public class PythonScalarFunctionRunnerTest extends AbstractPythonScalarFunction
PythonEnv pythonEnv,
RowType inputType, RowType outputType,
JobBundleFactory jobBundleFactory,
- String tempDir) {
- super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDir);
+ String[] tempDirs) {
+ super(taskName, resultReceiver, scalarFunctions, pythonEnv, inputType, outputType, tempDirs);
this.jobBundleFactory = jobBundleFactory;
}
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperatorTest.java b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperatorTest.java
new file mode 100644
index 0000000..35288cb
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/BaseRowPythonScalarFunctionOperatorTest.java
@@ -0,0 +1,103 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.api.common.typeinfo.Types;
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.table.dataformat.BaseRow;
+import org.apache.flink.table.dataformat.util.BaseRowUtil;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.util.BaseRowHarnessAssertor;
+import org.apache.flink.table.types.logical.RowType;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+import java.util.Collection;
+
+import static org.apache.flink.table.runtime.util.StreamRecordUtils.baserow;
+import static org.apache.flink.table.runtime.util.StreamRecordUtils.binaryrow;
+
+/**
+ * Tests for {@link BaseRowPythonScalarFunctionOperator}.
+ */
+public class BaseRowPythonScalarFunctionOperatorTest
+ extends PythonScalarFunctionOperatorTestBase<BaseRow, BaseRow, BaseRow, BaseRow> {
+
+ private final BaseRowHarnessAssertor assertor = new BaseRowHarnessAssertor(new TypeInformation[]{
+ Types.STRING,
+ Types.STRING,
+ Types.LONG
+ });
+
+ @Override
+ public AbstractPythonScalarFunctionOperator<BaseRow, BaseRow, BaseRow, BaseRow> getTestOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ return new PassThroughPythonScalarFunctionOperator(
+ scalarFunctions,
+ inputType,
+ outputType,
+ udfInputOffsets,
+ forwardedFieldCnt
+ );
+ }
+
+ @Override
+ public BaseRow newRow(boolean accumulateMsg, Object... fields) {
+ if (accumulateMsg) {
+ return baserow(fields);
+ } else {
+ return BaseRowUtil.setRetract(baserow(fields));
+ }
+ }
+
+ @Override
+ public void assertOutputEquals(String message, Collection<Object> expected, Collection<Object> actual) {
+ assertor.assertOutputEquals(message, expected, actual);
+ }
+
+ private static class PassThroughPythonScalarFunctionOperator extends BaseRowPythonScalarFunctionOperator {
+
+ PassThroughPythonScalarFunctionOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ super(scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFieldCnt);
+ }
+
+ @Override
+ public PythonFunctionRunner<BaseRow> createPythonFunctionRunner(
+ FnDataReceiver<BaseRow> resultReceiver) {
+ return new PassThroughPythonFunctionRunner<BaseRow>(resultReceiver) {
+ @Override
+ public BaseRow copy(BaseRow element) {
+ BaseRow row = binaryrow(element.getLong(0));
+ row.setHeader(element.getHeader());
+ return row;
+ }
+ };
+ }
+ }
+}
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PassThroughPythonFunctionRunner.java b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PassThroughPythonFunctionRunner.java
new file mode 100644
index 0000000..9660ab0
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PassThroughPythonFunctionRunner.java
@@ -0,0 +1,75 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.util.Preconditions;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A {@link PythonFunctionRunner} that just emit each input element.
+ *
+ * @param <IN> Type of the input elements.
+ */
+public abstract class PassThroughPythonFunctionRunner<IN> implements PythonFunctionRunner<IN> {
+
+ private boolean bundleStarted;
+ private final List<IN> bufferedElements;
+ private final FnDataReceiver<IN> resultReceiver;
+
+ PassThroughPythonFunctionRunner(FnDataReceiver<IN> resultReceiver) {
+ this.resultReceiver = Preconditions.checkNotNull(resultReceiver);
+ bundleStarted = false;
+ bufferedElements = new ArrayList<>();
+ }
+
+ @Override
+ public void open() {}
+
+ @Override
+ public void close() {}
+
+ @Override
+ public void startBundle() {
+ Preconditions.checkState(!bundleStarted);
+ bundleStarted = true;
+ }
+
+ @Override
+ public void finishBundle() throws Exception {
+ Preconditions.checkState(bundleStarted);
+ bundleStarted = false;
+
+ for (IN element : bufferedElements) {
+ resultReceiver.accept(element);
+ }
+ bufferedElements.clear();
+ }
+
+ @Override
+ public void processElement(IN element) {
+ bufferedElements.add(copy(element));
+ }
+
+ public abstract IN copy(IN element);
+}
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTest.java b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTest.java
new file mode 100644
index 0000000..d9a854e
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTest.java
@@ -0,0 +1,81 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.python.PythonFunctionRunner;
+import org.apache.flink.streaming.util.TestHarnessUtil;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.runtime.types.CRow;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.types.Row;
+
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
+
+import java.util.Collection;
+import java.util.Queue;
+
+/**
+ * Tests for {@link PythonScalarFunctionOperator}.
+ */
+public class PythonScalarFunctionOperatorTest extends PythonScalarFunctionOperatorTestBase<CRow, CRow, Row, Row> {
+
+ @Override
+ public AbstractPythonScalarFunctionOperator<CRow, CRow, Row, Row> getTestOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ return new PassThroughPythonScalarFunctionOperator(
+ scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFieldCnt);
+ }
+
+ @Override
+ public CRow newRow(boolean accumulateMsg, Object... fields) {
+ return new CRow(Row.of(fields), accumulateMsg);
+ }
+
+ @Override
+ public void assertOutputEquals(String message, Collection<Object> expected, Collection<Object> actual) {
+ TestHarnessUtil.assertOutputEquals(message, (Queue<Object>) expected, (Queue<Object>) actual);
+ }
+
+ private static class PassThroughPythonScalarFunctionOperator extends PythonScalarFunctionOperator {
+
+ PassThroughPythonScalarFunctionOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt) {
+ super(scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFieldCnt);
+ }
+
+ @Override
+ public PythonFunctionRunner<Row> createPythonFunctionRunner(
+ FnDataReceiver<Row> resultReceiver) {
+ return new PassThroughPythonFunctionRunner<Row>(resultReceiver) {
+ @Override
+ public Row copy(Row element) {
+ return Row.copy(element);
+ }
+ };
+ }
+ }
+}
diff --git a/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTestBase.java b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTestBase.java
new file mode 100644
index 0000000..a4c52a8
--- /dev/null
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/operators/python/PythonScalarFunctionOperatorTestBase.java
@@ -0,0 +1,223 @@
+/*
+ * 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.operators.python;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.python.PythonOptions;
+import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
+import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
+import org.apache.flink.table.functions.python.AbstractPythonScalarFunctionRunnerTest;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.types.logical.BigIntType;
+import org.apache.flink.table.types.logical.RowType;
+import org.apache.flink.table.types.logical.VarCharType;
+
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.concurrent.ConcurrentLinkedQueue;
+
+/**
+ * Base class for Python scalar function operator test. These test that:
+ *
+ * <ul>
+ * <li>Retraction flag is correctly forwarded to the downstream</li>
+ * <li>FinishBundle is called when checkpoint is encountered</li>
+ * <li>Watermarks are buffered and only sent to downstream when finishedBundle is triggered</li>
+ * </ul>
+ *
+ * @param <IN> Type of the input elements.
+ * @param <OUT> Type of the output elements.
+ * @param <UDFIN> Type of the UDF input type.
+ * @param <UDFOUT> Type of the UDF input type.
+ */
+public abstract class PythonScalarFunctionOperatorTestBase<IN, OUT, UDFIN, UDFOUT> {
+
+ @Test
+ public void testRetractionFieldKept() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+ testHarness.processElement(new StreamRecord<>(newRow(false, "c3", "c4", 1L), initialTime + 2));
+ testHarness.processElement(new StreamRecord<>(newRow(false, "c5", "c6", 2L), initialTime + 3));
+ testHarness.close();
+
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+ expectedOutput.add(new StreamRecord<>(newRow(false, "c3", "c4", 1L)));
+ expectedOutput.add(new StreamRecord<>(newRow(false, "c5", "c6", 2L)));
+
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+ }
+
+ @Test
+ public void testFinishBundleTriggeredOnCheckpoint() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+
+ Configuration conf = new Configuration();
+ conf.setInteger(PythonOptions.MAX_BUNDLE_SIZE, 10);
+ testHarness.getExecutionConfig().setGlobalJobParameters(conf);
+
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+
+ // checkpoint trigger finishBundle
+ testHarness.prepareSnapshotPreBarrier(0L);
+
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+
+ testHarness.close();
+ }
+
+ @Test
+ public void testFinishBundleTriggeredByCount() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+ Configuration conf = new Configuration();
+ conf.setInteger(PythonOptions.MAX_BUNDLE_SIZE, 2);
+ testHarness.getExecutionConfig().setGlobalJobParameters(conf);
+
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+ assertOutputEquals("FinishBundle should not be triggered.", expectedOutput, testHarness.getOutput());
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 1L), initialTime + 2));
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 1L)));
+
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+
+ testHarness.close();
+ }
+
+ @Test
+ public void testFinishBundleTriggeredByTime() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+ Configuration conf = new Configuration();
+ conf.setInteger(PythonOptions.MAX_BUNDLE_SIZE, 10);
+ conf.setLong(PythonOptions.MAX_BUNDLE_TIME_MILLS, 1000L);
+ testHarness.getExecutionConfig().setGlobalJobParameters(conf);
+
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+ assertOutputEquals("FinishBundle should not be triggered.", expectedOutput, testHarness.getOutput());
+
+ testHarness.setProcessingTime(1000L);
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+
+ testHarness.close();
+ }
+
+ @Test
+ public void testFinishBundleTriggeredByClose() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+ Configuration conf = new Configuration();
+ conf.setInteger(PythonOptions.MAX_BUNDLE_SIZE, 10);
+ testHarness.getExecutionConfig().setGlobalJobParameters(conf);
+
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+ assertOutputEquals("FinishBundle should not be triggered.", expectedOutput, testHarness.getOutput());
+
+ testHarness.close();
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+ }
+
+ @Test
+ public void testWatermarkProcessedOnFinishBundle() throws Exception {
+ OneInputStreamOperatorTestHarness<IN, OUT> testHarness = getTestHarness();
+ Configuration conf = new Configuration();
+ conf.setInteger(PythonOptions.MAX_BUNDLE_SIZE, 10);
+ testHarness.getExecutionConfig().setGlobalJobParameters(conf);
+ long initialTime = 0L;
+ ConcurrentLinkedQueue<Object> expectedOutput = new ConcurrentLinkedQueue<>();
+
+ testHarness.open();
+
+ testHarness.processElement(new StreamRecord<>(newRow(true, "c1", "c2", 0L), initialTime + 1));
+ testHarness.processWatermark(initialTime + 2);
+ assertOutputEquals("Watermark has been processed", expectedOutput, testHarness.getOutput());
+
+ // checkpoint trigger finishBundle
+ testHarness.prepareSnapshotPreBarrier(0L);
+
+ expectedOutput.add(new StreamRecord<>(newRow(true, "c1", "c2", 0L)));
+ expectedOutput.add(new Watermark(initialTime + 2));
+
+ assertOutputEquals("Output was not correct.", expectedOutput, testHarness.getOutput());
+
+ testHarness.close();
+ }
+
+ private OneInputStreamOperatorTestHarness<IN, OUT> getTestHarness() throws Exception {
+ RowType dataType = new RowType(Arrays.asList(
+ new RowType.RowField("f1", new VarCharType()),
+ new RowType.RowField("f2", new VarCharType()),
+ new RowType.RowField("f3", new BigIntType())));
+ AbstractPythonScalarFunctionOperator<IN, OUT, UDFIN, UDFOUT> operator = getTestOperator(
+ new PythonFunctionInfo[] {
+ new PythonFunctionInfo(
+ AbstractPythonScalarFunctionRunnerTest.DummyPythonFunction.INSTANCE,
+ new Integer[]{0})
+ },
+ dataType,
+ dataType,
+ new int[]{2},
+ 2
+ );
+
+ return new OneInputStreamOperatorTestHarness<>(operator);
+ }
+
+ public abstract AbstractPythonScalarFunctionOperator<IN, OUT, UDFIN, UDFOUT> getTestOperator(
+ PythonFunctionInfo[] scalarFunctions,
+ RowType inputType,
+ RowType outputType,
+ int[] udfInputOffsets,
+ int forwardedFieldCnt);
+
+ public abstract IN newRow(boolean accumulateMsg, Object... fields);
+
+ public abstract void assertOutputEquals(String message, Collection<Object> expected, Collection<Object> actual);
+}
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/BaseRowCoderTest.java b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoderTest.java
similarity index 98%
rename from flink-python/src/test/java/org/apache/flink/table/functions/python/coders/BaseRowCoderTest.java
rename to flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoderTest.java
index 57eb79b..7b5e281 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/BaseRowCoderTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/BaseRowCoderTest.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.table.dataformat.BaseRow;
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/CoderTestBase.java b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/CoderTestBase.java
similarity index 98%
rename from flink-python/src/test/java/org/apache/flink/table/functions/python/coders/CoderTestBase.java
rename to flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/CoderTestBase.java
index 642a544..308fe5d 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/CoderTestBase.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/CoderTestBase.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.testutils.CustomEqualityMatcher;
import org.apache.flink.testutils.DeeplyEqualsChecker;
diff --git a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/RowCoderTest.java b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/RowCoderTest.java
similarity index 96%
rename from flink-python/src/test/java/org/apache/flink/table/functions/python/coders/RowCoderTest.java
rename to flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/RowCoderTest.java
index 93b1c00..e8b6483 100644
--- a/flink-python/src/test/java/org/apache/flink/table/functions/python/coders/RowCoderTest.java
+++ b/flink-python/src/test/java/org/apache/flink/table/runtime/typeutils/coders/RowCoderTest.java
@@ -16,7 +16,7 @@
* limitations under the License.
*/
-package org.apache.flink.table.functions.python.coders;
+package org.apache.flink.table.runtime.typeutils.coders;
import org.apache.flink.types.Row;
diff --git a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/AbstractStreamOperatorTestHarness.java b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/AbstractStreamOperatorTestHarness.java
index c033951..8c7713a 100644
--- a/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/AbstractStreamOperatorTestHarness.java
+++ b/flink-streaming-java/src/test/java/org/apache/flink/streaming/util/AbstractStreamOperatorTestHarness.java
@@ -545,6 +545,13 @@ public class AbstractStreamOperatorTestHarness<OUT> implements AutoCloseable {
}
/**
+ * Calls {@link StreamOperator#prepareSnapshotPreBarrier(long)}.
+ */
+ public void prepareSnapshotPreBarrier(long checkpointId) throws Exception {
+ operator.prepareSnapshotPreBarrier(checkpointId);
+ }
+
+ /**
* Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions, org.apache.flink.runtime.state.CheckpointStreamFactory)}.
*/
public OperatorSubtaskState snapshot(long checkpointId, long timestamp) throws Exception {