Merge pull request #2289 from ClickHouse/add-column-names

Add support for partial insert

Author: Paultagoras

client-v2/src/main/java/com/clickhouse/client/api/Client.java

@@ -1425,9 +1425,23 @@ public CompletableFuture<InsertResponse> insert(String tableName, List<?> data,
      * @return {@code CompletableFuture<InsertResponse>} - a promise to insert response
      */
     public CompletableFuture<InsertResponse> insert(String tableName, InputStream data, ClickHouseFormat format) {
-        return insert(tableName, data, format, new InsertSettings());
+        return insert(tableName, Collections.emptyList(), data, format, new InsertSettings());
     }
 
+    /**
+     * <p>Sends write request to database. Input data is read from the input stream.</p>
+     *
+     * @param tableName - destination table name
+     * @param data  - data stream to insert
+     * @param format - format of the data in the stream
+     * @return {@code CompletableFuture<InsertResponse>} - a promise to insert response
+     */
+    public CompletableFuture<InsertResponse> insert(String tableName, List<String> columnNames, InputStream data, ClickHouseFormat format) {
+        return insert(tableName, columnNames, data, format, new InsertSettings());
+    }
+
+
+
     /**
      * Sends write request to database. Input data is read from the input stream.
      *
@@ -1441,6 +1455,23 @@ public CompletableFuture<InsertResponse> insert(String tableName,
                                                     InputStream data,
                                                     ClickHouseFormat format,
                                                     InsertSettings settings) {
+        return insert(tableName, Collections.emptyList(), data, format, settings);
+    }
+    /**
+     * Sends write request to database. Input data is read from the input stream.
+     *
+     * @param tableName - destination table name
+     * @param columnNames - list of column names to insert data into. If null or empty, all columns will be used.
+     * @param data  - data stream to insert
+     * @param format - format of the data in the stream
+     * @param settings - insert operation settings
+     * @return {@code CompletableFuture<InsertResponse>} - a promise to insert response
+     */
+    public CompletableFuture<InsertResponse> insert(String tableName,
+                                                    List<String> columnNames,
+                                                    InputStream data,
+                                                    ClickHouseFormat format,
+                                                    InsertSettings settings) {
 
         final int writeBufferSize = settings.getInputStreamCopyBufferSize() <= 0 ?
                 Integer.parseInt(configuration.getOrDefault(ClientConfigProperties.CLIENT_NETWORK_BUFFER_SIZE.getKey(),
@@ -1451,7 +1482,7 @@ public CompletableFuture<InsertResponse> insert(String tableName,
             throw new IllegalArgumentException("Buffer size must be greater than 0");
         }
 
-        return insert(tableName, new DataStreamWriter() {
+        return insert(tableName, columnNames, new DataStreamWriter() {
                     @Override
                     public void onOutput(OutputStream out) throws IOException {
                         byte[] buffer = new byte[writeBufferSize];
@@ -1480,11 +1511,29 @@ public void onRetry() throws IOException {
      * @return {@code CompletableFuture<InsertResponse>} - a promise to insert response
      */
     public CompletableFuture<InsertResponse> insert(String tableName,
+                                                    DataStreamWriter writer,
+                                                    ClickHouseFormat format,
+                                                    InsertSettings settings) {
+        return insert(tableName, Collections.emptyList(), writer, format, settings);
+    }
+
+    /**
+     * Does an insert request to a server. Data is pushed when a {@link DataStreamWriter#onOutput(OutputStream)} is called.
+     *
+     * @param tableName - target table name
+     * @param columnNames - list of column names to insert data into. If null or empty, all columns will be used.
+     * @param writer - {@link DataStreamWriter} implementation
+     * @param format - source format
+     * @param settings - operation settings
+     * @return {@code CompletableFuture<InsertResponse>} - a promise to insert response
+     */
+    public CompletableFuture<InsertResponse> insert(String tableName,
+                                     List<String> columnNames,
                                      DataStreamWriter writer,
                                      ClickHouseFormat format,
                                      InsertSettings settings) {
 
-        String operationId = (String) settings.getOperationId();
+        String operationId = settings.getOperationId();
         ClientStatisticsHolder clientStats = null;
         if (operationId != null) {
             clientStats = globalClientStats.remove(operationId);
@@ -1509,8 +1558,18 @@ public CompletableFuture<InsertResponse> insert(String tableName,
 
         settings.setOption(ClientConfigProperties.INPUT_OUTPUT_FORMAT.getKey(), format.name());
         final InsertSettings finalSettings = settings;
-        final String sqlStmt = "INSERT INTO " + tableName + " FORMAT " + format.name();
-        finalSettings.serverSetting(ClickHouseHttpProto.QPARAM_QUERY_STMT, sqlStmt);
+
+        StringBuilder sqlStmt = new StringBuilder("INSERT INTO ").append(tableName);
+        if (columnNames != null && !columnNames.isEmpty()) {
+            sqlStmt.append(" (");
+            for (String columnName : columnNames) {
+                sqlStmt.append(columnName).append(", ");
+            }
+            sqlStmt.deleteCharAt(sqlStmt.length() - 2);
+            sqlStmt.append(")");
+        }
+        sqlStmt.append(" FORMAT ").append(format.name());
+        finalSettings.serverSetting(ClickHouseHttpProto.QPARAM_QUERY_STMT, sqlStmt.toString());
         responseSupplier = () -> {
             long startTime = System.nanoTime();
             // Selecting some node

client-v2/src/test/java/com/clickhouse/client/insert/InsertTests.java

@@ -294,6 +294,31 @@ public void insertRawDataSimple(String tableName) throws Exception {
         assertEquals((int)response.getWrittenRows(), numberOfRecords );
     }
 
+    @Test(groups = { "integration" })
+    public void insertRawDataFewerColumns() throws Exception {
+        final String tableName = "raw_data_select_columns_table";
+        final String createSQL = "CREATE TABLE " + tableName +
+                " (Id UInt32, event_ts Timestamp, name String, p1 Int64, p2 String, p3 String, p4 Int8) ENGINE = MergeTree() ORDER BY ()";
+        List<String> columnNames = Arrays.asList("Id", "event_ts", "name", "p1", "p2");
+
+        initTable(tableName, createSQL);
+
+        settings.setInputStreamCopyBufferSize(8198 * 2);
+        ByteArrayOutputStream data = new ByteArrayOutputStream();
+        PrintWriter writer = new PrintWriter(data);
+        for (int i = 0; i < 1000; i++) {
+            writer.printf("%d\t%s\t%s\t%d\t%s\n", i, "2021-01-01 00:00:00", "name" + i, i, "p2");
+        }
+        writer.flush();
+        InsertResponse response = client.insert(tableName, columnNames, new ByteArrayInputStream(data.toByteArray()),
+                ClickHouseFormat.TSV, settings).get(30, TimeUnit.SECONDS);
+        OperationMetrics metrics = response.getMetrics();
+        assertEquals((int)response.getWrittenRows(), 1000 );
+
+        List<GenericRecord> records = client.queryAll("SELECT * FROM " + tableName);
+        assertEquals(records.size(), 1000);
+    }
+
     @DataProvider(name = "insertRawDataSimpleDataProvider")
     public static Object[][] insertRawDataSimpleDataProvider() {
         return new Object[][] {