Optimize MergeUsing.buildColumnListFromOnCondition()

h2/src/main/org/h2/command/dml/MergeUsing.java

@@ -342,9 +342,7 @@ public class MergeUsing extends Prepared {
         onCondition.mapColumns(targetTableFilter, 1);
 
         if (keys == null) {
-            HashSet<Column> targetColumns = buildColumnListFromOnCondition(
-                    targetTableFilter);
-            keys = targetColumns.toArray(new Column[0]);
+            keys = buildColumnListFromOnCondition(targetTableFilter.getTable());
         }
         if (keys.length == 0) {
             throw DbException.get(ErrorCode.COLUMN_NOT_FOUND_1,
@@ -410,19 +408,11 @@ public class MergeUsing extends Prepared {
         targetMatchQuery.prepare();
     }
 
-    private HashSet<Column> buildColumnListFromOnCondition(
-            TableFilter anyTableFilter) {
-        HashSet<Column> filteredColumns = new HashSet<>();
+    private Column[] buildColumnListFromOnCondition(Table table) {
         HashSet<Column> columns = new HashSet<>();
-        ExpressionVisitor visitor = ExpressionVisitor
-                .getColumnsVisitor(columns);
+        ExpressionVisitor visitor = ExpressionVisitor.getColumnsVisitor(columns, table);
         onCondition.isEverything(visitor);
-        for (Column c : columns) {
-            if (c != null && c.getTable() == anyTableFilter.getTable()) {
-                filteredColumns.add(c);
-            }
-        }
-        return filteredColumns;
+        return columns.toArray(new Column[0]);
     }
 
     private Expression appendOnCondition(Update updateCommand) {

h2/src/main/org/h2/constraint/ConstraintCheck.java

@@ -6,7 +6,6 @@
 package org.h2.constraint;
 
 import java.util.HashSet;
-import java.util.Iterator;
 import org.h2.api.ErrorCode;
 import org.h2.engine.Session;
 import org.h2.expression.Expression;
@@ -122,12 +121,7 @@ public class ConstraintCheck extends Constraint {
     @Override
     public HashSet<Column> getReferencedColumns(Table table) {
         HashSet<Column> columns = new HashSet<>();
-        expr.isEverything(ExpressionVisitor.getColumnsVisitor(columns));
-        for (Iterator<Column> it = columns.iterator(); it.hasNext();) {
-            if (it.next().getTable() != table) {
-                it.remove();
-            }
-        }
+        expr.isEverything(ExpressionVisitor.getColumnsVisitor(columns, table));
         return columns;
     }
 

h2/src/main/org/h2/expression/ExpressionVisitor.java

@@ -239,10 +239,11 @@ public class ExpressionVisitor {
      * Create a new visitor to get all referenced columns.
      *
      * @param columns the columns map
+     * @param table table to gather columns from, or {@code null} to gather all columns
      * @return the new visitor
      */
-    public static ExpressionVisitor getColumnsVisitor(HashSet<Column> columns) {
-        return new ExpressionVisitor(GET_COLUMNS2, 0, null, null, null, null, null, columns);
+    public static ExpressionVisitor getColumnsVisitor(HashSet<Column> columns, Table table) {
+        return new ExpressionVisitor(GET_COLUMNS2, 0, null, null, table, null, null, columns);
     }
 
     public static ExpressionVisitor getMaxModificationIdVisitor() {
@@ -269,8 +270,11 @@ public class ExpressionVisitor {
     void addColumn1(Column column) {
         columns1.add(column);
     }
+
     void addColumn2(Column column) {
-        columns2.add(column);
+        if (table == null || table == column.getTable()) {
+            columns2.add(column);
+        }
     }
 
     /**