Skip to content

Commit dfb161b

Browse files
chernserchernser
authored
Merge pull request #1816 from ClickHouse/clientv2_buffer_optimization
[client-v2] reader buffer optimization
2 parents f3a4ab8 + ba490d5 commit dfb161b

File tree

18 files changed

+497
-232
lines changed

18 files changed

+497
-232
lines changed

‎.github/workflows/analysis.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,6 @@ jobs:
9595
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # Needed to get PR information, if any
9696
SONAR_TOKEN: ${{ secrets.SONAR_TOKEN }}
9797
run: |
98-
mvn --batch-mode -DclickhouseVersion=$PREFERRED_LTS_VERSION \
98+
mvn --batch-mode -DclickhouseVersion=$PREFERRED_LTS_VERSION -Dclient.tests.useNewImplementation=true \
9999
-Panalysis verify org.sonarsource.scanner.maven:sonar-maven-plugin:sonar
100100
continue-on-error: true

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

Lines changed: 44 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import com.clickhouse.client.api.data_formats.RowBinaryFormatReader;
1212
import com.clickhouse.client.api.data_formats.RowBinaryWithNamesAndTypesFormatReader;
1313
import com.clickhouse.client.api.data_formats.RowBinaryWithNamesFormatReader;
14+
import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
1415
import com.clickhouse.client.api.data_formats.internal.MapBackedRecord;
1516
import com.clickhouse.client.api.data_formats.internal.ProcessParser;
1617
import com.clickhouse.client.api.enums.Protocol;
@@ -20,6 +21,7 @@
2021
import com.clickhouse.client.api.insert.InsertSettings;
2122
import com.clickhouse.client.api.insert.POJOSerializer;
2223
import com.clickhouse.client.api.insert.SerializerNotFoundException;
24+
import com.clickhouse.client.api.internal.BasicObjectsPool;
2325
import com.clickhouse.client.api.internal.ClickHouseLZ4OutputStream;
2426
import com.clickhouse.client.api.internal.ClientStatisticsHolder;
2527
import com.clickhouse.client.api.internal.ClientV1AdaptorHelper;
@@ -68,6 +70,7 @@
6870
import java.util.HashMap;
6971
import java.util.HashSet;
7072
import java.util.LinkedHashMap;
73+
import java.util.LinkedList;
7174
import java.util.List;
7275
import java.util.Map;
7376
import java.util.Set;
@@ -753,6 +756,22 @@ public Builder setMaxRetries(int maxRetries) {
753756
return this;
754757
}
755758

759+
/**
760+
* Configures client to reuse allocated byte buffers for numbers. It affects how binary format reader is working.
761+
* If set to 'true' then {@link Client#newBinaryFormatReader(QueryResponse)} will construct reader that will
762+
* reuse buffers for numbers. It improves performance for large datasets by reducing number of allocations
763+
* (therefore GC pressure).
764+
* Enabling this feature is safe because each reader suppose to be used by a single thread and readers are not reused.
765+
*
766+
* Default is false.
767+
* @param reuse - if to reuse buffers
768+
* @return
769+
*/
770+
public Builder allowBinaryReaderToReuseBuffers(boolean reuse) {
771+
this.configuration.put("client_allow_binary_reader_to_reuse_buffers", String.valueOf(reuse));
772+
return this;
773+
}
774+
756775
public Client build() {
757776
setDefaults();
758777

@@ -866,6 +885,10 @@ private void setDefaults() {
866885
if (!configuration.containsKey(ClickHouseClientOption.RETRY.getKey())) {
867886
setMaxRetries(3);
868887
}
888+
889+
if (!configuration.containsKey("client_allow_binary_reader_to_reuse_buffers")) {
890+
allowBinaryReaderToReuseBuffers(false);
891+
}
869892
}
870893
}
871894

@@ -1442,10 +1465,10 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
14421465
settings.setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes);
14431466
settings.waitEndOfQuery(true); // we rely on the summery
14441467

1445-
final QuerySettings finalSettings = settings;
14461468
return query(sqlQuery, settings).thenApply(response -> {
14471469
try {
1448-
return new Records(response, finalSettings);
1470+
1471+
return new Records(response, newBinaryFormatReader(response));
14491472
} catch (Exception e) {
14501473
throw new ClientException("Failed to get query response", e);
14511474
}
@@ -1462,13 +1485,14 @@ public CompletableFuture<Records> queryRecords(String sqlQuery, QuerySettings se
14621485
public List<GenericRecord> queryAll(String sqlQuery) {
14631486
try {
14641487
int operationTimeout = getOperationTimeout();
1465-
QuerySettings settings = new QuerySettings().waitEndOfQuery(true);
1488+
QuerySettings settings = new QuerySettings().setFormat(ClickHouseFormat.RowBinaryWithNamesAndTypes)
1489+
.waitEndOfQuery(true);
14661490
try (QueryResponse response = operationTimeout == 0 ? query(sqlQuery, settings).get() :
14671491
query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
14681492
List<GenericRecord> records = new ArrayList<>();
14691493
if (response.getResultRows() > 0) {
14701494
RowBinaryWithNamesAndTypesFormatReader reader =
1471-
new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream(), response.getSettings());
1495+
(RowBinaryWithNamesAndTypesFormatReader) newBinaryFormatReader(response);
14721496

14731497
Map<String, Object> record;
14741498
while (reader.readRecord((record = new LinkedHashMap<>()))) {
@@ -1569,28 +1593,38 @@ public CompletableFuture<CommandResponse> execute(String sql) {
15691593
* @param schema
15701594
* @return
15711595
*/
1572-
public static ClickHouseBinaryFormatReader newBinaryFormatReader(QueryResponse response, TableSchema schema) {
1596+
public ClickHouseBinaryFormatReader newBinaryFormatReader(QueryResponse response, TableSchema schema) {
15731597
ClickHouseBinaryFormatReader reader = null;
1598+
// Using caching buffer allocator is risky so this parameter is not exposed to the user
1599+
boolean useCachingBufferAllocator = MapUtils.getFlag(configuration, "client_allow_binary_reader_to_reuse_buffers");
1600+
BinaryStreamReader.ByteBufferAllocator byteBufferPool = useCachingBufferAllocator ?
1601+
new BinaryStreamReader.CachingByteBufferAllocator() :
1602+
new BinaryStreamReader.DefaultByteBufferAllocator();
1603+
15741604
switch (response.getFormat()) {
15751605
case Native:
1576-
reader = new NativeFormatReader(response.getInputStream(), response.getSettings());
1606+
reader = new NativeFormatReader(response.getInputStream(), response.getSettings(),
1607+
byteBufferPool);
15771608
break;
15781609
case RowBinaryWithNamesAndTypes:
1579-
reader = new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream(), response.getSettings());
1610+
reader = new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream(), response.getSettings(),
1611+
byteBufferPool);
15801612
break;
15811613
case RowBinaryWithNames:
1582-
reader = new RowBinaryWithNamesFormatReader(response.getInputStream(), response.getSettings(), schema);
1614+
reader = new RowBinaryWithNamesFormatReader(response.getInputStream(), response.getSettings(), schema,
1615+
byteBufferPool);
15831616
break;
15841617
case RowBinary:
1585-
reader = new RowBinaryFormatReader(response.getInputStream(), response.getSettings(), schema);
1618+
reader = new RowBinaryFormatReader(response.getInputStream(), response.getSettings(), schema,
1619+
byteBufferPool);
15861620
break;
15871621
default:
15881622
throw new IllegalArgumentException("Unsupported format: " + response.getFormat());
15891623
}
15901624
return reader;
15911625
}
15921626

1593-
public static ClickHouseBinaryFormatReader newBinaryFormatReader(QueryResponse response) {
1627+
public ClickHouseBinaryFormatReader newBinaryFormatReader(QueryResponse response) {
15941628
return newBinaryFormatReader(response, null);
15951629
}
15961630

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/ClickHouseBinaryFormatReader.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import java.util.Map;
2020
import java.util.UUID;
2121

22-
public interface ClickHouseBinaryFormatReader {
22+
public interface ClickHouseBinaryFormatReader extends AutoCloseable {
2323

2424
/**
2525
* Reads a single value from the stream.

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/NativeFormatReader.java‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,9 @@ public class NativeFormatReader extends AbstractBinaryFormatReader {
2323

2424
private int blockRowIndex;
2525

26-
public NativeFormatReader(InputStream inputStream) {
27-
this(inputStream, null);
28-
}
29-
30-
public NativeFormatReader(InputStream inputStream, QuerySettings settings) {
31-
super(inputStream, settings, null);
26+
public NativeFormatReader(InputStream inputStream, QuerySettings settings,
27+
BinaryStreamReader.ByteBufferAllocator byteBufferAllocator) {
28+
super(inputStream, settings, null, byteBufferAllocator);
3229
readNextRecord();
3330
}
3431

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryFormatReader.java‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.clickhouse.client.api.data_formats;
22

33
import com.clickhouse.client.api.data_formats.internal.AbstractBinaryFormatReader;
4+
import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
45
import com.clickhouse.client.api.metadata.TableSchema;
56
import com.clickhouse.client.api.query.QuerySettings;
67
import com.clickhouse.data.ClickHouseColumn;
@@ -12,12 +13,9 @@
1213

1314
public class RowBinaryFormatReader extends AbstractBinaryFormatReader {
1415

15-
public RowBinaryFormatReader(InputStream inputStream, TableSchema schema) {
16-
this(inputStream, null, schema);
17-
}
18-
19-
public RowBinaryFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema) {
20-
super(inputStream, querySettings, schema);
16+
public RowBinaryFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema,
17+
BinaryStreamReader.ByteBufferAllocator byteBufferAllocator) {
18+
super(inputStream, querySettings, schema, byteBufferAllocator);
2119
readNextRecord();
2220
}
2321

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesAndTypesFormatReader.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
66
import com.clickhouse.client.api.metadata.TableSchema;
77
import com.clickhouse.client.api.query.QuerySettings;
8-
import com.clickhouse.data.ClickHouseColumn;
98

109
import java.io.EOFException;
1110
import java.io.IOException;
@@ -17,8 +16,9 @@
1716

1817
public class RowBinaryWithNamesAndTypesFormatReader extends AbstractBinaryFormatReader implements Iterator<Map<String, Object>> {
1918

20-
public RowBinaryWithNamesAndTypesFormatReader(InputStream inputStream, QuerySettings querySettings) {
21-
super(inputStream, querySettings, null);
19+
public RowBinaryWithNamesAndTypesFormatReader(InputStream inputStream, QuerySettings querySettings,
20+
BinaryStreamReader.ByteBufferAllocator byteBufferAllocator) {
21+
super(inputStream, querySettings, null, byteBufferAllocator);
2222
readSchema();
2323
}
2424

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/RowBinaryWithNamesFormatReader.java‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,27 +4,21 @@
44
import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
55
import com.clickhouse.client.api.metadata.TableSchema;
66
import com.clickhouse.client.api.query.QuerySettings;
7-
import com.clickhouse.data.ClickHouseColumn;
87

98
import java.io.EOFException;
109
import java.io.IOException;
1110
import java.io.InputStream;
1211
import java.util.ArrayList;
1312
import java.util.Collections;
1413
import java.util.List;
15-
import java.util.Map;
1614

1715
public class RowBinaryWithNamesFormatReader extends AbstractBinaryFormatReader {
1816

1917
private List<String> columns = null;
2018

21-
public RowBinaryWithNamesFormatReader(InputStream inputStream, TableSchema schema) {
22-
this(inputStream, null, schema);
23-
readNextRecord();
24-
}
25-
26-
public RowBinaryWithNamesFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema) {
27-
super(inputStream, querySettings, schema);
19+
public RowBinaryWithNamesFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema,
20+
BinaryStreamReader.ByteBufferAllocator byteBufferAllocator) {
21+
super(inputStream, querySettings, schema, byteBufferAllocator);
2822
int nCol = 0;
2923
try {
3024
nCol = BinaryStreamReader.readVarInt(input);

‎client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.clickhouse.client.api.ClientException;
44
import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
5+
import com.clickhouse.client.api.internal.BasicObjectsPool;
56
import com.clickhouse.client.api.internal.MapUtils;
67
import com.clickhouse.client.api.metadata.TableSchema;
78
import com.clickhouse.client.api.query.NullValueException;
@@ -41,6 +42,7 @@
4142
import java.util.concurrent.ConcurrentHashMap;
4243
import java.util.concurrent.atomic.AtomicBoolean;
4344
import java.util.function.Function;
45+
import java.util.function.Supplier;
4446

4547
public abstract class AbstractBinaryFormatReader implements ClickHouseBinaryFormatReader {
4648

@@ -58,7 +60,8 @@ public abstract class AbstractBinaryFormatReader implements ClickHouseBinaryForm
5860

5961
private volatile boolean hasNext = true;
6062

61-
protected AbstractBinaryFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema) {
63+
protected AbstractBinaryFormatReader(InputStream inputStream, QuerySettings querySettings, TableSchema schema,
64+
BinaryStreamReader.ByteBufferAllocator byteBufferAllocator) {
6265
this.input = inputStream;
6366
this.settings = querySettings == null ? Collections.emptyMap() : new HashMap<>(querySettings.getAllSettings());
6467
boolean useServerTimeZone = (boolean) this.settings.get(ClickHouseClientOption.USE_SERVER_TIME_ZONE.getKey());
@@ -67,7 +70,7 @@ protected AbstractBinaryFormatReader(InputStream inputStream, QuerySettings quer
6770
if (timeZone == null) {
6871
throw new ClientException("Time zone is not set. (useServerTimezone:" + useServerTimeZone + ")");
6972
}
70-
this.binaryStreamReader = new BinaryStreamReader(inputStream, timeZone, LOG);
73+
this.binaryStreamReader = new BinaryStreamReader(inputStream, timeZone, LOG, byteBufferAllocator);
7174
setSchema(schema);
7275
}
7376

@@ -133,12 +136,12 @@ protected void readNextRecord() {
133136
try {
134137
nextRecordEmpty.set(true);
135138
if (!readRecord(nextRecord)) {
136-
hasNext = false;
139+
endReached();
137140
} else {
138141
nextRecordEmpty.compareAndSet(true, false);
139142
}
140143
} catch (IOException e) {
141-
hasNext = false;
144+
endReached();
142145
throw new ClientException("Failed to read next row", e);
143146
}
144147
}
@@ -165,7 +168,7 @@ public Map<String, Object> next() {
165168
return null;
166169
}
167170
} catch (IOException e) {
168-
hasNext = false;
171+
endReached();
169172
throw new ClientException("Failed to read row", e);
170173
}
171174
}
@@ -621,4 +624,9 @@ public LocalDateTime getLocalDateTime(int index) {
621624
}
622625
return (LocalDateTime) value;
623626
}
627+
628+
@Override
629+
public void close() throws Exception {
630+
input.close();
631+
}
624632
}

0 commit comments

Comments
 (0)