fixed deserialization of ClickHouseBitmap

chernser · chernser · commit daedbda7dd71 · 2024-10-22T13:07:37.000-07:00
diff --git a/clickhouse-data/src/main/java/com/clickhouse/data/value/ClickHouseBitmap.java b/clickhouse-data/src/main/java/com/clickhouse/data/value/ClickHouseBitmap.java
@@ -354,7 +354,7 @@ public static ClickHouseBitmap deserialize(DataInputStream in, ClickHouseDataTyp
             byte[] bytes = new byte[2 + byteLen * cardinality];
             bytes[0] = (byte) flag;
             bytes[1] = cardinality;
-            in.read(bytes, 2, bytes.length - 2);
+            in.readFully(bytes, 2, bytes.length - 2);
 
             rb = ClickHouseBitmap.deserialize(bytes, innerType);
         } else {
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/ClickHouseBinaryFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/ClickHouseBinaryFormatReader.java
@@ -1,6 +1,7 @@
 package com.clickhouse.client.api.data_formats;
 
 import com.clickhouse.client.api.metadata.TableSchema;
+import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
 import com.clickhouse.data.value.ClickHouseGeoPolygonValue;
@@ -519,4 +520,8 @@ public interface ClickHouseBinaryFormatReader extends AutoCloseable {
     LocalDateTime getLocalDateTime(int index);
 
     TableSchema getSchema();
+
+    ClickHouseBitmap getClickHouseBitmap(String colName);
+
+    ClickHouseBitmap getClickHouseBitmap(int index);
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/AbstractBinaryFormatReader.java
@@ -8,6 +8,7 @@
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.client.config.ClickHouseClientOption;
 import com.clickhouse.data.ClickHouseColumn;
+import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
 import com.clickhouse.data.value.ClickHouseGeoPolygonValue;
@@ -661,6 +662,16 @@ public LocalDateTime getLocalDateTime(int index) {
         return (LocalDateTime) value;
     }
 
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(String colName) {
+        return readValue(colName);
+    }
+
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(int index) {
+        return readValue(index);
+    }
+
     @Override
     public void close() throws Exception {
         input.close();
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryReaderBackedRecord.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryReaderBackedRecord.java
@@ -2,6 +2,7 @@
 
 import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
 import com.clickhouse.client.api.query.GenericRecord;
+import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
 import com.clickhouse.data.value.ClickHouseGeoPolygonValue;
@@ -356,4 +357,14 @@ public Object getObject(String colName) {
     public Object getObject(int index) {
         return reader.readValue(index);
     }
+
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(String colName) {
+        return reader.readValue(colName);
+    }
+
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(int index) {
+        return reader.readValue(index);
+    }
 }
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryStreamReader.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/BinaryStreamReader.java
@@ -216,7 +216,7 @@ public <T> T readValue(ClickHouseColumn column, Class<?> typeHint) throws IOExce
                     return null;
 //                case SimpleAggregateFunction:
                 case AggregateFunction:
-                    return (T) ClickHouseBitmap.deserialize(input, column.getNestedColumns().get(0).getDataType());
+                    return (T) readBitmap( column);
                 default:
                     throw new IllegalArgumentException("Unsupported data type: " + column.getDataType());
             }
@@ -885,6 +885,11 @@ public static boolean isReadToPrimitive(ClickHouseDataType dataType) {
                 return false;
         }
     }
+
+    private ClickHouseBitmap readBitmap(ClickHouseColumn column) throws IOException {
+        return ClickHouseBitmap.deserialize(input, column.getNestedColumns().get(0).getDataType());
+    }
+
     /**
      * Byte allocator that caches preallocated byte arrays for small sizes.
      */
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java b/client-v2/src/main/java/com/clickhouse/client/api/data_formats/internal/MapBackedRecord.java
@@ -6,6 +6,7 @@
 import com.clickhouse.client.api.query.NullValueException;
 import com.clickhouse.data.ClickHouseColumn;
 import com.clickhouse.data.value.ClickHouseArrayValue;
+import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
 import com.clickhouse.data.value.ClickHouseGeoPolygonValue;
@@ -482,6 +483,16 @@ public LocalDateTime getLocalDateTime(int index) {
         return (LocalDateTime) value;
     }
 
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(String colName) {
+        return readValue(colName);
+    }
+
+    @Override
+    public ClickHouseBitmap getClickHouseBitmap(int index) {
+        return readValue(index);
+    }
+
     @Override
     public Object getObject(String colName) {
         return readValue(colName);
diff --git a/client-v2/src/main/java/com/clickhouse/client/api/query/GenericRecord.java b/client-v2/src/main/java/com/clickhouse/client/api/query/GenericRecord.java
@@ -1,5 +1,6 @@
 package com.clickhouse.client.api.query;
 
+import com.clickhouse.data.value.ClickHouseBitmap;
 import com.clickhouse.data.value.ClickHouseGeoMultiPolygonValue;
 import com.clickhouse.data.value.ClickHouseGeoPointValue;
 import com.clickhouse.data.value.ClickHouseGeoPolygonValue;
@@ -489,4 +490,8 @@ public interface GenericRecord {
     Object getObject(String colName);
 
     Object getObject(int index);
+
+    ClickHouseBitmap getClickHouseBitmap(String colName);
+
+    ClickHouseBitmap getClickHouseBitmap(int index);
 }
diff --git a/client-v2/src/test/java/com/clickhouse/client/insert/SamplePOJO.java b/client-v2/src/test/java/com/clickhouse/client/insert/SamplePOJO.java
@@ -680,10 +680,10 @@ public static String generateTableCreateSQL(String tableName) {
                 "array Array(String), " +
                 "tuple Tuple(UInt64, Int32, String), " +
                 "map Map(String, Int32), " +
-                "nested Nested (innerInt Int32, innerString String, innerNullableInt Nullable(Int32))" +
-//                "groupBitmapUint32 AggregateFunction(groupBitmap, UInt32)," +
-                // TODO: fix this
-//                "groupBitmapUint64 AggregateFunction(groupBitmap, UInt64)" +
+                "nested Nested (innerInt Int32, innerString String, " +
+                "innerNullableInt Nullable(Int32)), " +
+                "groupBitmapUint32 AggregateFunction(groupBitmap, UInt32), " +
+                "groupBitmapUint64 AggregateFunction(groupBitmap, UInt64) " +
                 ") ENGINE = Memory";
     }
 }
diff --git a/client-v2/src/test/java/com/clickhouse/client/query/AggregateFuncDTO.java b/client-v2/src/test/java/com/clickhouse/client/query/AggregateFuncDTO.java
@@ -0,0 +1,54 @@
+package com.clickhouse.client.query;
+
+import com.clickhouse.data.value.ClickHouseBitmap;
+
+import java.util.Objects;
+import java.util.Random;
+
+public class AggregateFuncDTO {
+
+    private ClickHouseBitmap groupBitmapUint32;
+    private ClickHouseBitmap groupBitmapUint64;
+
+    public AggregateFuncDTO() {
+        Random random = new Random();
+        this.groupBitmapUint32 = ClickHouseBitmap.wrap(random.ints(5, Integer.MAX_VALUE - 100, Integer.MAX_VALUE).toArray());
+        this.groupBitmapUint64 = ClickHouseBitmap.wrap(random.longs(5, Long.MAX_VALUE - 100, Long.MAX_VALUE).toArray());
+    }
+
+    public ClickHouseBitmap getGroupBitmapUint32() {
+        return groupBitmapUint32;
+    }
+
+    public void setGroupBitmapUint32(ClickHouseBitmap groupBitmapUint32) {
+        this.groupBitmapUint32 = groupBitmapUint32;
+    }
+
+    public ClickHouseBitmap getGroupBitmapUint64() {
+        return groupBitmapUint64;
+    }
+
+    public void setGroupBitmapUint64(ClickHouseBitmap groupBitmapUint64) {
+        this.groupBitmapUint64 = groupBitmapUint64;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        AggregateFuncDTO that = (AggregateFuncDTO) o;
+        return Objects.equals(groupBitmapUint32, that.groupBitmapUint32) && Objects.equals(groupBitmapUint64, that.groupBitmapUint64);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(groupBitmapUint32, groupBitmapUint64);
+    }
+
+    public static String generateTableCreateSQL(String tableName) {
+        return "CREATE TABLE " + tableName + " (" +
+                "groupBitmapUint32 AggregateFunction(groupBitmap, UInt32), " +
+                "groupBitmapUint64 AggregateFunction(groupBitmap, UInt64) " +
+                ") ENGINE = MergeTree() ORDER BY tuple()";
+    }
+}
diff --git a/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java b/client-v2/src/test/java/com/clickhouse/client/query/QueryTests.java
@@ -16,6 +16,7 @@
 import com.clickhouse.client.api.ServerException;
 import com.clickhouse.client.api.data_formats.ClickHouseBinaryFormatReader;
 import com.clickhouse.client.api.enums.Protocol;
+import com.clickhouse.client.api.insert.InsertResponse;
 import com.clickhouse.client.api.insert.InsertSettings;
 import com.clickhouse.client.api.metadata.TableSchema;
 import com.clickhouse.client.api.metrics.ClientMetrics;
@@ -26,6 +27,7 @@
 import com.clickhouse.client.api.query.QueryResponse;
 import com.clickhouse.client.api.query.QuerySettings;
 import com.clickhouse.client.api.query.Records;
+import com.clickhouse.client.insert.SamplePOJO;
 import com.clickhouse.data.ClickHouseDataType;
 import com.clickhouse.data.ClickHouseFormat;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -1618,6 +1620,36 @@ public static BigDecimal cropDecimal(BigDecimal value, int scale) {
         return new BigDecimal(bi, scale);
     }
 
+    @Test(groups = { "integration" }, enabled = true)
+    public void testReadingBitmap() throws Exception {
+        final String tableName = "bitmaps_test_table";
+        final String createSQL = AggregateFuncDTO.generateTableCreateSQL(tableName);
+        final AggregateFuncDTO pojo = new AggregateFuncDTO();
+
+        try {
+            client.execute("DROP TABLE IF EXISTS " + tableName).get();
+            client.execute(createSQL).get();
+        } catch (Exception e) {
+            throw e;
+        }
+
+        client.register(AggregateFuncDTO.class, client.getTableSchema(tableName, "default"));
+
+        try (InsertResponse response = client.insert(tableName, Collections.singletonList(pojo)).get(30, TimeUnit.SECONDS)) {
+            Assert.assertEquals(response.getWrittenRows(), 1);
+        }
+
+        try (QueryResponse queryResponse =
+                     client.query("SELECT * FROM " + tableName + " LIMIT 1").get(30, TimeUnit.SECONDS)) {
+
+            ClickHouseBinaryFormatReader reader = client.newBinaryFormatReader(queryResponse);
+            Assert.assertNotNull(reader.next());
+            Assert.assertFalse(reader.hasNext());
+
+            Assert.assertEquals(reader.getClickHouseBitmap("groupBitmapUint32"), pojo.getGroupBitmapUint32());
+            Assert.assertEquals(reader.getClickHouseBitmap("groupBitmapUint64"), pojo.getGroupBitmapUint64());
+        }
+    }
 
     protected Client.Builder newClient() {
         ClickHouseNode node = getServer(ClickHouseProtocol.HTTP);
