Improve JDBC insert performance (#2165)

* First naive implementation.

* Refine implementation

* Enable testDateTypes

* remove jdbc v1 from tests

* Use StatementType instead of boolean

Author: mzitnik

jdbc-v2/src/main/java/com/clickhouse/jdbc/PreparedStatementImpl.java

@@ -38,9 +38,11 @@
 import java.time.format.DateTimeFormatter;
 import java.time.format.DateTimeFormatterBuilder;
 import java.time.temporal.ChronoField;
+import java.util.ArrayList;
 import java.util.Calendar;
 import java.util.Collection;
 import java.util.GregorianCalendar;
+import java.util.List;
 import java.util.Map;
 
 public class PreparedStatementImpl extends StatementImpl implements PreparedStatement, JdbcV2Wrapper {
@@ -57,11 +59,19 @@ public class PreparedStatementImpl extends StatementImpl implements PreparedStat
     String originalSql;
     String [] sqlSegments;
     Object [] parameters;
+    String insertIntoSQL;
+
+    StatementType statementType;
     public PreparedStatementImpl(ConnectionImpl connection, String sql) throws SQLException {
         super(connection);
-        this.originalSql = sql;
+        this.originalSql = sql.trim();
         //Split the sql string into an array of strings around question mark tokens
-        this.sqlSegments = sql.split("\\?");
+        this.sqlSegments = originalSql.split("\\?");
+        this.statementType = parseStatementType(originalSql);
+
+        if (statementType == StatementType.INSERT) {
+            insertIntoSQL = originalSql.substring(0, originalSql.indexOf("VALUES") + 6);
+        }
 
         //Create an array of objects to store the parameters
         if (originalSql.contains("?")) {
@@ -86,6 +96,19 @@ private String compileSql() {
         return sb.toString();
     }
 
+    private String valuesSql() {
+        StringBuilder sb = new StringBuilder("(");
+        for (int i = 0; i < parameters.length; i++) {
+            if (i > 0) {
+                sb.append(", ");
+            }
+            sb.append(parameters[i]);
+        }
+        sb.append(")");
+        LOG.trace("Compiled Value SQL: {}", sb);
+        return sb.toString();
+    }
+
     @Override
     public ResultSet executeQuery() throws SQLException {
         checkClosed();
@@ -228,7 +251,33 @@ public boolean execute() throws SQLException {
     @Override
     public void addBatch() throws SQLException {
         checkClosed();
-        addBatch(compileSql());
+        if (statementType == StatementType.INSERT) {
+            addBatch(valuesSql());
+        } else {
+            addBatch(compileSql());
+        }
+    }
+
+    @Override
+    public int[] executeBatch() throws SQLException {
+        checkClosed();
+        if (statementType == StatementType.INSERT && !batch.isEmpty()) {
+            List<Integer> results = new ArrayList<>();
+            // write insert into as batch to avoid multiple requests
+            StringBuilder sb = new StringBuilder();
+            sb.append(insertIntoSQL).append(" ");
+            for (String sql : batch) {
+                sb.append(sql).append(",");
+            }
+            sb.setCharAt(sb.length() - 1, ';');
+            results.add(executeUpdate(sb.toString()));
+            // clear batch and re-add insert into
+            batch.clear();
+            return results.stream().mapToInt(i -> i).toArray();
+        } else {
+            // run executeBatch
+            return super.executeBatch();
+        }
     }
 
     @Override

jdbc-v2/src/main/java/com/clickhouse/jdbc/StatementImpl.java

@@ -33,12 +33,11 @@ public class StatementImpl implements Statement, JdbcV2Wrapper {
     protected boolean closed;
     private ResultSetImpl currentResultSet;
     private OperationMetrics metrics;
-    private List<String> batch;
+    protected List<String> batch;
     private String lastSql;
     private volatile String lastQueryId;
     private String schema;
     private int maxRows;
-
     public StatementImpl(ConnectionImpl connection) throws SQLException {
         this.connection = connection;
         this.queryTimeout = 0;
@@ -460,11 +459,9 @@ public void clearBatch() throws SQLException {
     public int[] executeBatch() throws SQLException {
         checkClosed();
         List<Integer> results = new ArrayList<>();
-
-        for(String sql : batch) {
+        for (String sql : batch) {
             results.add(executeUpdate(sql));
         }
-
         return results.stream().mapToInt(i -> i).toArray();
     }
 

jdbc-v2/src/test/java/com/clickhouse/jdbc/PreparedStatementTest.java

@@ -8,6 +8,7 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
+import java.sql.Statement;
 import java.sql.Timestamp;
 import java.sql.Types;
 import java.time.ZoneId;
@@ -276,4 +277,40 @@ void testWithClause() throws Exception {
         }
         assertEquals(count, 100);
     }
+
+    @Test(groups = { "integration" })
+    void testInsert() throws Exception {
+        int ROWS = 100000;
+        try (Connection conn = getJdbcConnection()) {
+            for (int j = 0; j < 10; j++) {
+                try (Statement stmt = conn.createStatement()) {
+                    stmt.execute("CREATE TABLE insert_batch (`id` UInt32, `name` String) ENGINE = Memory");
+                }
+                String insertQuery = "INSERT INTO insert_batch (id, name) VALUES (?,?)";
+                try (PreparedStatement stmt = conn.prepareStatement(insertQuery)) {
+                    for (int i = 0; i < ROWS; i++) {
+                        stmt.setInt(1, i);
+                        stmt.setString(2, "name" + i);
+                        stmt.addBatch();
+                    }
+                    long startBatchTime = System.currentTimeMillis();
+                    stmt.executeBatch();
+                    long endBatchTime = System.currentTimeMillis();
+                    System.out.println("Insertion Time for Final Batch: " + (endBatchTime - startBatchTime) + " ms");
+                }
+                // count rows
+                try (Statement stmt = conn.createStatement()) {
+                    try (ResultSet rs = stmt.executeQuery("SELECT count(*) FROM insert_batch")) {
+                        assertTrue(rs.next());
+                        assertEquals(rs.getInt(1), ROWS);
+                    }
+                }
+                try (Statement stmt = conn.createStatement()) {
+                    stmt.execute("DROP TABLE insert_batch");
+                }
+
+            }
+
+        }
+    }
 }