Updates based on JingsongLi'comment

Author: beyond1920

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/join/JoinConditionWithNullFilters.java

@@ -18,16 +18,16 @@
 
 package org.apache.flink.table.runtime.operators.join;
 
-import org.apache.flink.api.common.functions.AbstractRichFunction;
-import org.apache.flink.configuration.Configuration;
+import org.apache.flink.api.java.operators.translation.WrappingFunction;
+import org.apache.flink.streaming.api.operators.KeyContext;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.data.binary.NullAwareGetters;
 import org.apache.flink.table.runtime.generated.JoinCondition;
 
 /** Utility to take null filters into consideration when apply join condition. */
-public class JoinConditionWithNullFilters extends AbstractRichFunction implements JoinCondition {
+public class JoinConditionWithNullFilters extends WrappingFunction<JoinCondition> implements JoinCondition {
 
-    private final JoinCondition backingJoinCondition;
+    private static final long serialVersionUID = 1L;
 
     /** Should filter null keys. */
     private final int[] nullFilterKeys;
@@ -38,42 +38,28 @@ public class JoinConditionWithNullFilters extends AbstractRichFunction implement
     /** Filter null to all keys. */
     private final boolean filterAllNulls;
 
-    private NullAwareGetters joinKey;
+    private KeyContext keyContext;
 
     public JoinConditionWithNullFilters(
-            JoinCondition backingJoinCondition, boolean[] filterNullKeys) {
-        this.backingJoinCondition = backingJoinCondition;
+            JoinCondition backingJoinCondition, boolean[] filterNullKeys, KeyContext keyContext) {
+        super(backingJoinCondition);
         this.nullFilterKeys = NullAwareJoinHelper.getNullFilterKeys(filterNullKeys);
         this.nullSafe = nullFilterKeys.length == 0;
         this.filterAllNulls = nullFilterKeys.length == filterNullKeys.length;
-    }
-
-    @Override
-    public void open(Configuration parameters) throws Exception {
-        super.open(parameters);
-        backingJoinCondition.open(parameters);
-    }
-
-    @Override
-    public void close() throws Exception {
-        super.close();
-        backingJoinCondition.close();
+        this.keyContext = keyContext;
     }
 
     @Override
     public boolean apply(RowData left, RowData right) {
         if (!nullSafe) { // is not null safe, return false if any null exists
             // key is always BinaryRowData
+            NullAwareGetters joinKey = (NullAwareGetters) keyContext.getCurrentKey();
             if (filterAllNulls ? joinKey.anyNull() : joinKey.anyNull(nullFilterKeys)) {
                 // find null present, return false directly
                 return false;
             }
         }
         // test condition
-        return backingJoinCondition.apply(left, right);
-    }
-
-    public void setCurrentJoinKey(NullAwareGetters joinKey) {
-        this.joinKey = joinKey;
+        return wrappedFunction.apply(left, right);
     }
 }

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/join/stream/AbstractStreamingJoinOperator.java

@@ -85,11 +85,10 @@ public AbstractStreamingJoinOperator(
     @Override
     public void open() throws Exception {
         super.open();
-
         JoinCondition condition =
                 generatedJoinCondition.newInstance(getRuntimeContext().getUserCodeClassLoader());
-        condition.setRuntimeContext(getRuntimeContext());
-        this.joinCondition = new JoinConditionWithNullFilters(condition, filterNullKeys);
+        this.joinCondition = new JoinConditionWithNullFilters(condition, filterNullKeys, this);
+        this.joinCondition.setRuntimeContext(getRuntimeContext());
         this.joinCondition.open(new Configuration());
 
         this.collector = new TimestampedCollector<>(output);

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/join/stream/StreamingJoinOperator.java

@@ -21,7 +21,6 @@
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.table.data.GenericRowData;
 import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.data.binary.NullAwareGetters;
 import org.apache.flink.table.data.util.RowDataUtil;
 import org.apache.flink.table.data.utils.JoinedRowData;
 import org.apache.flink.table.runtime.generated.GeneratedJoinCondition;
@@ -208,7 +207,6 @@ private void processElement(
         RowKind inputRowKind = input.getRowKind();
         input.setRowKind(RowKind.INSERT); // erase RowKind for later state updating
 
-        joinCondition.setCurrentJoinKey((NullAwareGetters) getCurrentKey());
         AssociatedRecords associatedRecords =
                 AssociatedRecords.of(input, inputIsLeft, otherSideStateView, joinCondition);
         if (isAccumulateMsg) { // record is accumulate

flink-table/flink-table-runtime-blink/src/main/java/org/apache/flink/table/runtime/operators/join/stream/StreamingSemiAntiJoinOperator.java

@@ -20,7 +20,6 @@
 
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.data.binary.NullAwareGetters;
 import org.apache.flink.table.data.util.RowDataUtil;
 import org.apache.flink.table.runtime.generated.GeneratedJoinCondition;
 import org.apache.flink.table.runtime.operators.join.stream.state.JoinInputSideSpec;
@@ -103,7 +102,6 @@ public void open() throws Exception {
     @Override
     public void processElement1(StreamRecord<RowData> element) throws Exception {
         RowData input = element.getValue();
-        joinCondition.setCurrentJoinKey((NullAwareGetters) getCurrentKey());
         AssociatedRecords associatedRecords =
                 AssociatedRecords.of(input, true, rightRecordStateView, joinCondition);
         if (associatedRecords.isEmpty()) {
@@ -168,7 +166,6 @@ public void processElement2(StreamRecord<RowData> element) throws Exception {
         RowKind inputRowKind = input.getRowKind();
         input.setRowKind(RowKind.INSERT); // erase RowKind for later state updating
 
-        joinCondition.setCurrentJoinKey((NullAwareGetters) getCurrentKey());
         AssociatedRecords associatedRecords =
                 AssociatedRecords.of(input, false, leftRecordStateView, joinCondition);
         if (isAccumulateMsg) { // record is accumulate