Merge pull request #1814 from javiercj93/bitmap-serialize

[client-v2] Including support for bitmap serialization
ClickHouse · Sep 12, 2024 · 28de7cd · 28de7cd
2 parents f9b0f7a + 2539301
commit 28de7cd
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 6 deletions.
diff --git a/client-v2/pom.xml b/client-v2/pom.xml
@@ -131,6 +131,11 @@
             <version>2.35.2</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.roaringbitmap</groupId>
+            <artifactId>RoaringBitmap</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/client-v2/src/main/java/com/clickhouse/client/api/internal/SerializerUtils.java b/client-v2/src/main/java/com/clickhouse/client/api/internal/SerializerUtils.java
@@ -4,9 +4,12 @@
 import com.clickhouse.client.api.ClientException;
 import com.clickhouse.client.api.data_formats.internal.BinaryStreamReader;
 import com.clickhouse.client.api.query.POJOSetter;
+import com.clickhouse.data.ClickHouseAggregateFunction;
 import com.clickhouse.data.ClickHouseColumn;
 import com.clickhouse.data.ClickHouseDataType;
 import com.clickhouse.data.format.BinaryStreamUtils;
+import com.clickhouse.data.value.ClickHouseBitmap;
+
 import org.objectweb.asm.ClassWriter;
 import org.objectweb.asm.MethodVisitor;
 import org.objectweb.asm.Opcodes;
@@ -60,6 +63,9 @@ public static void serializeData(OutputStream stream, Object value, ClickHouseCo
             case Map:
                 serializeMapData(stream, value, column);
                 break;
+            case AggregateFunction:
+                serializeAggregateFunction(stream, value, column);
+                break;
             default:
                 serializePrimitiveData(stream, value, column);
                 break;
@@ -194,6 +200,13 @@ private static void serializePrimitiveData(OutputStream stream, Object value, Cl
         }
     }
 
+    private static void serializeAggregateFunction(OutputStream stream, Object value, ClickHouseColumn column) throws IOException {
+        if (column.getAggregateFunction() == ClickHouseAggregateFunction.groupBitmap) {
+            BinaryStreamUtils.writeBitmap(stream, (ClickHouseBitmap) value);
+        } else {
+            throw new UnsupportedOperationException("Unsupported aggregate function: " + column.getAggregateFunction());
+        }
+    }
 
     public static Integer convertToInteger(Object value) {
         if (value instanceof Integer) {

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
@@ -1,8 +1,7 @@
 package com.clickhouse.client.insert;
 
-import com.clickhouse.client.api.metadata.TableSchema;
-import com.clickhouse.data.ClickHouseColumn;
-import com.clickhouse.data.ClickHouseEnum;
+import com.clickhouse.data.value.ClickHouseBitmap;
+
 import org.apache.commons.lang3.RandomStringUtils;
 
 import java.math.BigDecimal;
@@ -14,13 +13,11 @@
 import java.time.LocalDateTime;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.UUID;
-import java.util.stream.IntStream;
 
 public class SamplePOJO {
     private int int8;
@@ -76,6 +73,9 @@ public class SamplePOJO {
     private List<Integer> nestedInnerInt;
     private List<String> nestedInnerString;
 
+    private ClickHouseBitmap groupBitmapUint32;
+    private ClickHouseBitmap groupBitmapUint64;
+
     public SamplePOJO() {
         final Random random = new Random();
         int8 = random.nextInt(128);
@@ -154,6 +154,9 @@ public SamplePOJO() {
         List<String> innerString = new ArrayList<>();
         innerString.add(RandomStringUtils.randomAlphabetic(1, 256));
         nestedInnerString = innerString;
+
+        groupBitmapUint32 = ClickHouseBitmap.wrap(random.ints(5, Integer.MAX_VALUE - 100, Integer.MAX_VALUE).toArray());
+        groupBitmapUint64 = ClickHouseBitmap.wrap(random.longs(5, Long.MAX_VALUE - 100, Long.MAX_VALUE).toArray());
     }
 
     public int getInt8() {
@@ -460,6 +463,22 @@ public void setNestedInnerString(List<String> nestedInnerString) {
         this.nestedInnerString = nestedInnerString;
     }
 
+    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;
+    }
+
     public static String generateTableCreateSQL(String tableName) {
         return "CREATE TABLE " + tableName + " (" +
                 "int8 Int8, " +
@@ -501,7 +520,9 @@ public static String generateTableCreateSQL(String tableName) {
                 "array Array(String), " +
                 "tuple Tuple(Int32, Int32, Int32, Int32, Int32, Int32, Int32, Int32, Int32, Int32), " +
                 "map Map(String, Int32), " +
-                "nested Nested (innerInt Int32, innerString String)" +
+                "nested Nested (innerInt Int32, innerString String)," +
+                "groupBitmapUint32 AggregateFunction(groupBitmap, UInt32)," +
+                "groupBitmapUint64 AggregateFunction(groupBitmap, UInt64)" +
                 ") ENGINE = Memory";
     }
 }