[FLINK-33365] include filters with Lookup joins

Signed-off-by: David Radley <[email protected]>
Co-authored-by: Benchao Li <[email protected]>
Co-authored-by: Sergey Nuyanzin <[email protected]>

Author: 3 people

flink-connector-jdbc/src/main/java/org/apache/flink/connector/jdbc/statement/FieldNamedPreparedStatement.java

@@ -69,7 +69,18 @@ public interface FieldNamedPreparedStatement extends AutoCloseable {
      */
     static FieldNamedPreparedStatement prepareStatement(
             Connection connection, String sql, String[] fieldNames) throws SQLException {
-        return FieldNamedPreparedStatementImpl.prepareStatement(connection, sql, fieldNames);
+        return FieldNamedPreparedStatementImpl.prepareStatement(connection, sql, fieldNames, "", 0);
+    }
+
+    static FieldNamedPreparedStatement prepareStatement(
+            Connection connection,
+            String sql,
+            String[] fieldNames,
+            String additionalPredicates,
+            int numberOfDynamicParams)
+            throws SQLException {
+        return FieldNamedPreparedStatementImpl.prepareStatement(
+                connection, sql, fieldNames, additionalPredicates, numberOfDynamicParams);
     }
 
     /**

flink-connector-jdbc/src/main/java/org/apache/flink/connector/jdbc/statement/FieldNamedPreparedStatementImpl.java

@@ -178,26 +178,46 @@ public void close() throws SQLException {
     // ----------------------------------------------------------------------------------------
 
     public static FieldNamedPreparedStatement prepareStatement(
-            Connection connection, String sql, String[] fieldNames) throws SQLException {
+            Connection connection,
+            String sql,
+            String[] fieldNames,
+            String additionalPredicates,
+            int numberOfDynamicParams)
+            throws SQLException {
         checkNotNull(connection, "connection must not be null.");
         checkNotNull(sql, "sql must not be null.");
         checkNotNull(fieldNames, "fieldNames must not be null.");
 
         if (sql.contains("?")) {
             throw new IllegalArgumentException("SQL statement must not contain ? character.");
         }
+        sql = sql + additionalPredicates;
 
         HashMap<String, List<Integer>> parameterMap = new HashMap<>();
         String parsedSQL = parseNamedStatement(sql, parameterMap);
+
         // currently, the statements must contain all the field parameters
-        checkArgument(parameterMap.size() == fieldNames.length);
-        int[][] indexMapping = new int[fieldNames.length][];
-        for (int i = 0; i < fieldNames.length; i++) {
+        final int parameterMapSize = parameterMap.size();
+        final int fieldNamesLength = fieldNames.length;
+        checkArgument(
+                parameterMapSize == fieldNamesLength,
+                "Expected "
+                        + fieldNamesLength
+                        + " fields, but the parsing found "
+                        + parameterMapSize);
+        int[][] indexMapping = new int[fieldNamesLength + numberOfDynamicParams][];
+        int numberOfNameBasedParams = 0;
+        for (int i = 0; i < fieldNamesLength; i++) {
             String fieldName = fieldNames[i];
             checkArgument(
                     parameterMap.containsKey(fieldName),
                     fieldName + " doesn't exist in the parameters of SQL statement: " + sql);
             indexMapping[i] = parameterMap.get(fieldName).stream().mapToInt(v -> v).toArray();
+            numberOfNameBasedParams += parameterMap.get(fieldName).size();
+        }
+        for (int i = 0; i < numberOfDynamicParams; ++i) {
+            // FieldNamedPreparedStatement is 0-based, however, PreparedStatement is 1-based
+            indexMapping[i + fieldNamesLength] = new int[] {i + numberOfNameBasedParams + 1};
         }
 
         return new FieldNamedPreparedStatementImpl(

flink-connector-jdbc/src/main/java/org/apache/flink/connector/jdbc/table/JdbcDynamicTableSource.java

@@ -110,7 +110,9 @@ public LookupRuntimeProvider getLookupRuntimeProvider(LookupContext context) {
                         DataType.getFieldNames(physicalRowDataType).toArray(new String[0]),
                         DataType.getFieldDataTypes(physicalRowDataType).toArray(new DataType[0]),
                         keyNames,
-                        rowType);
+                        rowType,
+                        resolvedPredicates,
+                        pushdownParams);
         if (cache != null) {
             return PartialCachingLookupProvider.of(lookupFunction, cache);
         } else {

flink-connector-jdbc/src/main/java/org/apache/flink/connector/jdbc/table/JdbcRowDataLookupFunction.java

@@ -37,6 +37,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.io.Serializable;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -45,6 +46,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
@@ -63,6 +65,9 @@ public class JdbcRowDataLookupFunction extends LookupFunction {
     private final JdbcRowConverter jdbcRowConverter;
     private final JdbcRowConverter lookupKeyRowConverter;
 
+    private final List<String> resolvedPredicates;
+    private final Serializable[] pushdownParams;
+
     private transient FieldNamedPreparedStatement statement;
 
     public JdbcRowDataLookupFunction(
@@ -71,11 +76,15 @@ public JdbcRowDataLookupFunction(
             String[] fieldNames,
             DataType[] fieldTypes,
             String[] keyNames,
-            RowType rowType) {
+            RowType rowType,
+            List<String> resolvedPredicates,
+            Serializable[] pushdownParams) {
         checkNotNull(options, "No JdbcOptions supplied.");
         checkNotNull(fieldNames, "No fieldNames supplied.");
         checkNotNull(fieldTypes, "No fieldTypes supplied.");
         checkNotNull(keyNames, "No keyNames supplied.");
+        checkNotNull(resolvedPredicates, "No resolvedPredicates supplied.");
+        checkNotNull(pushdownParams, "No pushdownParams supplied.");
         this.connectionProvider = new SimpleJdbcConnectionProvider(options);
         this.keyNames = keyNames;
         List<String> nameList = Arrays.asList(fieldNames);
@@ -103,6 +112,8 @@ public JdbcRowDataLookupFunction(
                                 Arrays.stream(keyTypes)
                                         .map(DataType::getLogicalType)
                                         .toArray(LogicalType[]::new)));
+        this.resolvedPredicates = resolvedPredicates;
+        this.pushdownParams = pushdownParams;
     }
 
     @Override
@@ -116,6 +127,15 @@ public void open(FunctionContext context) throws Exception {
         }
     }
 
+    private FieldNamedPreparedStatement setPredicateParams(FieldNamedPreparedStatement statement)
+            throws SQLException {
+        for (int i = 0; i < pushdownParams.length; ++i) {
+            statement.setObject(i + keyNames.length, pushdownParams[i]);
+        }
+
+        return statement;
+    }
+
     /**
      * This is a lookup method which is called by Flink framework in runtime.
      *
@@ -127,6 +147,7 @@ public Collection<RowData> lookup(RowData keyRow) {
             try {
                 statement.clearParameters();
                 statement = lookupKeyRowConverter.toExternal(keyRow, statement);
+                statement = setPredicateParams(statement);
                 try (ResultSet resultSet = statement.executeQuery()) {
                     ArrayList<RowData> rows = new ArrayList<>();
                     while (resultSet.next()) {
@@ -167,7 +188,21 @@ public Collection<RowData> lookup(RowData keyRow) {
 
     private void establishConnectionAndStatement() throws SQLException, ClassNotFoundException {
         Connection dbConn = connectionProvider.getOrEstablishConnection();
-        statement = FieldNamedPreparedStatement.prepareStatement(dbConn, query, keyNames);
+        String additionalPredicates = "";
+        if (!resolvedPredicates.isEmpty()) {
+            String joinedConditions =
+                    resolvedPredicates.stream()
+                            .map(pred -> String.format("(%s)", pred))
+                            .collect(Collectors.joining(" AND "));
+            if (keyNames.length == 0) {
+                additionalPredicates = " WHERE " + joinedConditions;
+            } else {
+                additionalPredicates = " AND " + joinedConditions;
+            }
+        }
+        statement =
+                FieldNamedPreparedStatement.prepareStatement(
+                        dbConn, query, keyNames, additionalPredicates, pushdownParams.length);
     }
 
     @Override