[connector/flink] Fix NPE and improve project row code in Lookup

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/source/lookup/FlinkAsyncLookupFunction.java

@@ -24,6 +24,7 @@
 import com.alibaba.fluss.connector.flink.source.lookup.LookupNormalizer.RemainingFilter;
 import com.alibaba.fluss.connector.flink.utils.FlinkConversions;
 import com.alibaba.fluss.connector.flink.utils.FlinkRowToFlussRowConverter;
+import com.alibaba.fluss.connector.flink.utils.FlinkUtils;
 import com.alibaba.fluss.connector.flink.utils.FlussRowToFlinkRowConverter;
 import com.alibaba.fluss.exception.TableNotExistException;
 import com.alibaba.fluss.metadata.TablePath;
@@ -88,12 +89,17 @@ public void open(FunctionContext context) {
         // TODO: convert to Fluss GenericRow to avoid unnecessary deserialization
         flinkRowToFlussRowConverter =
                 FlinkRowToFlussRowConverter.create(
-                        FlinkLookupFunction.filterRowType(flinkRowType, pkIndexes),
+                        FlinkUtils.projectRowType(flinkRowType, pkIndexes),
                         table.getDescriptor().getKvFormat());
+
+        final RowType outputRowType;
+        if (projection == null) {
+            outputRowType = flinkRowType;
+        } else {
+            outputRowType = FlinkUtils.projectRowType(flinkRowType, projection);
+        }
         flussRowToFlinkRowConverter =
-                new FlussRowToFlinkRowConverter(
-                        FlinkConversions.toFlussRowType(
-                                FlinkLookupFunction.filterRowType(flinkRowType, projection)));
+                new FlussRowToFlinkRowConverter(FlinkConversions.toFlussRowType(outputRowType));
         LOG.info("end open.");
     }
 
@@ -164,8 +170,7 @@ private void fetchResult(
                             resultFuture.complete(Collections.emptyList());
                         } else {
                             RowData flinkRow =
-                                    flussRowToFlinkRowConverter.toFlinkRowData(
-                                            ProjectedRow.from(projection).replaceRow(row));
+                                    flussRowToFlinkRowConverter.toFlinkRowData(maybeProject(row));
                             if (remainingFilter != null && !remainingFilter.isMatch(flinkRow)) {
                                 resultFuture.complete(Collections.emptyList());
                             } else {
@@ -176,6 +181,14 @@ private void fetchResult(
                 });
     }
 
+    private InternalRow maybeProject(InternalRow row) {
+        if (projection == null) {
+            return row;
+        }
+        // should not reuse objects for async operations
+        return ProjectedRow.from(projection).replaceRow(row);
+    }
+
     @Override
     public void close() throws Exception {
         LOG.info("start close ...");

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/source/lookup/FlinkLookupFunction.java

@@ -22,16 +22,15 @@
 import com.alibaba.fluss.config.Configuration;
 import com.alibaba.fluss.connector.flink.utils.FlinkConversions;
 import com.alibaba.fluss.connector.flink.utils.FlinkRowToFlussRowConverter;
+import com.alibaba.fluss.connector.flink.utils.FlinkUtils;
 import com.alibaba.fluss.connector.flink.utils.FlussRowToFlinkRowConverter;
 import com.alibaba.fluss.metadata.TablePath;
 import com.alibaba.fluss.row.InternalRow;
 import com.alibaba.fluss.row.ProjectedRow;
 
 import org.apache.flink.table.data.RowData;
-import org.apache.flink.table.data.utils.ProjectedRowData;
 import org.apache.flink.table.functions.FunctionContext;
 import org.apache.flink.table.functions.LookupFunction;
-import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.RowType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -59,7 +58,7 @@ public class FlinkLookupFunction extends LookupFunction {
     private transient FlussRowToFlinkRowConverter flussRowToFlinkRowConverter;
     private transient Connection connection;
     private transient Table table;
-    private transient ProjectedRowData projectedRowData;
+    @Nullable private transient ProjectedRow projectedRow;
 
     public FlinkLookupFunction(
             Configuration flussConfig,
@@ -78,16 +77,6 @@ public FlinkLookupFunction(
         this.projection = projection;
     }
 
-    static RowType filterRowType(RowType rowType, int[] filterIndex) {
-        LogicalType[] types = new LogicalType[filterIndex.length];
-        String[] names = new String[filterIndex.length];
-        for (int i = 0; i < filterIndex.length; i++) {
-            types[i] = rowType.getTypeAt(filterIndex[i]);
-            names[i] = rowType.getFieldNames().get(filterIndex[i]);
-        }
-        return RowType.of(rowType.isNullable(), types, names);
-    }
-
     @Override
     public void open(FunctionContext context) {
         LOG.info("start open ...");
@@ -96,11 +85,21 @@ public void open(FunctionContext context) {
         // TODO: convert to Fluss GenericRow to avoid unnecessary deserialization
         flinkRowToFlussRowConverter =
                 FlinkRowToFlussRowConverter.create(
-                        filterRowType(flinkRowType, pkIndexes),
+                        FlinkUtils.projectRowType(flinkRowType, pkIndexes),
                         table.getDescriptor().getKvFormat());
+
+        final RowType outputRowType;
+        if (projection == null) {
+            outputRowType = flinkRowType;
+            projectedRow = null;
+        } else {
+            outputRowType = FlinkUtils.projectRowType(flinkRowType, projection);
+            // reuse the projected row
+            projectedRow = ProjectedRow.from(projection);
+        }
         flussRowToFlinkRowConverter =
-                new FlussRowToFlinkRowConverter(
-                        FlinkConversions.toFlussRowType(filterRowType(flinkRowType, projection)));
+                new FlussRowToFlinkRowConverter(FlinkConversions.toFlussRowType(outputRowType));
+
         LOG.info("end open.");
     }
 
@@ -125,8 +124,7 @@ public Collection<RowData> lookup(RowData keyRow) {
                 InternalRow row = table.lookup(flussKeyRow).get().getRow();
                 if (row != null) {
                     RowData flinkRow =
-                            flussRowToFlinkRowConverter.toFlinkRowData(
-                                    ProjectedRow.from(projection).replaceRow(row));
+                            flussRowToFlinkRowConverter.toFlinkRowData(maybeProject(row));
                     if (remainingFilter == null || remainingFilter.isMatch(flinkRow)) {
                         return Collections.singletonList(flinkRow);
                     } else {
@@ -153,6 +151,13 @@ public Collection<RowData> lookup(RowData keyRow) {
         return Collections.emptyList();
     }
 
+    private InternalRow maybeProject(InternalRow row) {
+        if (projectedRow == null) {
+            return row;
+        }
+        return projectedRow.replaceRow(row);
+    }
+
     @Override
     public void close() throws Exception {
         LOG.info("start close ...");

fluss-connectors/fluss-connector-flink/src/main/java/com/alibaba/fluss/connector/flink/utils/FlinkUtils.java

@@ -0,0 +1,42 @@
+/*
+ * 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 com.alibaba.fluss.connector.flink.utils;
+
+import org.apache.flink.table.types.logical.LogicalType;
+import org.apache.flink.table.types.logical.RowType;
+
+/** Utils for Flink classes. */
+public class FlinkUtils {
+
+    /**
+     * Returns projected {@link RowType} by given projection indexes over original {@link RowType}.
+     *
+     * @param rowType the original row type
+     * @param projection the projection indexes
+     */
+    public static RowType projectRowType(RowType rowType, int[] projection) {
+        LogicalType[] types = new LogicalType[projection.length];
+        String[] names = new String[projection.length];
+        for (int i = 0; i < projection.length; i++) {
+            types[i] = rowType.getTypeAt(projection[i]);
+            names[i] = rowType.getFieldNames().get(projection[i]);
+        }
+        return RowType.of(rowType.isNullable(), types, names);
+    }
+}