improved queryAll() to use less mem

Author: chernser

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

@@ -67,6 +67,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -1404,10 +1405,11 @@ public List<GenericRecord> queryAll(String sqlQuery) {
                     query(sqlQuery, settings).get(operationTimeout, TimeUnit.MILLISECONDS)) {
                 List<GenericRecord> records = new ArrayList<>();
                 if (response.getResultRows() > 0) {
-                    ClickHouseBinaryFormatReader reader =
+                    RowBinaryWithNamesAndTypesFormatReader reader =
                             new RowBinaryWithNamesAndTypesFormatReader(response.getInputStream(), response.getSettings());
+
                     Map<String, Object> record;
-                    while ((record = reader.next()) != null) {
+                    while (reader.readRecord((record = new LinkedHashMap<>()))) {
                         records.add(new MapBackedRecord(record, reader.getSchema()));
                     }
                 }

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

@@ -33,7 +33,7 @@ public NativeFormatReader(InputStream inputStream, QuerySettings settings) {
     }
 
     @Override
-    protected boolean readRecord(Map<String, Object> record) throws IOException {
+    public boolean readRecord(Map<String, Object> record) throws IOException {
         if (currentBlock == null || blockRowIndex >= currentBlock.getnRows()) {
             if (!readBlock()) {
                 return false;

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

@@ -30,6 +30,7 @@
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -38,6 +39,7 @@
 import java.util.TimeZone;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 
 public abstract class AbstractBinaryFormatReader implements ClickHouseBinaryFormatReader {
 
@@ -70,7 +72,18 @@ protected AbstractBinaryFormatReader(InputStream inputStream, QuerySettings quer
     protected Map<String, Object> nextRecord = new ConcurrentHashMap<>();
 
 
-    protected boolean readRecord(Map<String, Object> record) throws IOException {
+    /**
+     * It is still internal method and should be used with care.
+     * Usually this method is called to read next record into internal object and affects hasNext() method.
+     * So after calling this one:
+     * - hasNext(), next() should not be called
+     * - stream should be read with readRecord() method fully
+     *
+     * @param record
+     * @return
+     * @throws IOException
+     */
+    public boolean readRecord(Map<String, Object> record) throws IOException {
         boolean firstColumn = true;
         for (ClickHouseColumn column : getSchema().getColumns()) {
             try {

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

@@ -346,4 +346,14 @@ public LocalDateTime getLocalDateTime(String colName) {
     public LocalDateTime getLocalDateTime(int index) {
         return reader.getLocalDateTime(index);
     }
+
+    @Override
+    public Object getObject(String colName) {
+        return reader.readValue(colName);
+    }
+
+    @Override
+    public Object getObject(int index) {
+        return reader.readValue(index);
+    }
 }

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

@@ -480,4 +480,14 @@ public LocalDateTime getLocalDateTime(String colName) {
     public LocalDateTime getLocalDateTime(int index) {
         return readValue(index);
     }
+
+    @Override
+    public Object getObject(String colName) {
+        return readValue(colName);
+    }
+
+    @Override
+    public Object getObject(int index) {
+        return readValue(index);
+    }
 }

client-v2/src/main/java/com/clickhouse/client/api/query/GenericRecord.java

@@ -19,7 +19,6 @@
 
 public interface GenericRecord {
 
-
     /**
      * Reads column with name `colName` as a string.
      *
@@ -486,4 +485,8 @@ public interface GenericRecord {
     LocalDateTime getLocalDateTime(String colName);
 
     LocalDateTime getLocalDateTime(int index);
+
+    Object getObject(String colName);
+
+    Object getObject(int index);
 }

client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java

@@ -180,9 +180,19 @@ public void testReadRecordsNoResult() throws Exception {
 
     @Test(groups = {"integration"})
     public void testQueryAll() throws Exception {
-        prepareDataSet(DATASET_TABLE, DATASET_COLUMNS, DATASET_VALUE_GENERATORS, 10);
-        GenericRecord hostnameRecord = client.queryAll("SELECT hostname()").stream().findFirst().get();
-        Assert.assertNotNull(hostnameRecord);
+        List<Map<String, Object>> dataset = prepareDataSet(DATASET_TABLE, DATASET_COLUMNS, DATASET_VALUE_GENERATORS, 10);
+        List<GenericRecord> records = client.queryAll("SELECT * FROM " + DATASET_TABLE + " LIMIT " + dataset.size());
+        Assert.assertFalse(records.isEmpty());
+
+        int i = 0;
+        for (String colDefinition : DATASET_COLUMNS) {
+
+            String colName = colDefinition.split(" ")[0];
+            List<Object> colValues = records.stream().map(r -> r.getObject(colName)).toList();
+            Assert.assertEquals(colValues.size(), dataset.size());
+            List<Object> dataValue = dataset.stream().map(d -> d.get(colName)).toList();
+            Assert.assertEquals(colValues, dataValue, "Failed for column " + colName);
+        }
     }
 
     @Test(groups = {"integration"})