Support variant type

clickhouse-data/src/main/java/com/clickhouse/data/ClickHouseColumn.java

@@ -37,12 +37,16 @@
 
 import java.io.Serializable;
 import java.lang.reflect.Array;
+import java.math.BigInteger;
 import java.time.OffsetDateTime;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.TimeZone;
 
@@ -65,6 +69,7 @@ public final class ClickHouseColumn implements Serializable {
     private static final String KEYWORD_OBJECT = ClickHouseDataType.Object.name();
     private static final String KEYWORD_MAP = ClickHouseDataType.Map.name();
     private static final String KEYWORD_NESTED = ClickHouseDataType.Nested.name();
+    private static final String KEYWORD_VARIANT = ClickHouseDataType.Variant.name();
 
     private int columnCount;
     private int columnIndex;
@@ -92,6 +97,8 @@ public final class ClickHouseColumn implements Serializable {
 
     private ClickHouseValue template;
 
+    private Map<Class<?>, Integer> classToVariantOrdNumMap;
+
     private static ClickHouseColumn update(ClickHouseColumn column) {
         column.enumConstants = ClickHouseEnum.EMPTY;
         int size = column.parameters.size();
@@ -273,6 +280,9 @@ private static ClickHouseColumn update(ClickHouseColumn column) {
             case Nothing:
                 column.template = ClickHouseEmptyValue.INSTANCE;
                 break;
+            case Variant:
+                column.template = ClickHouseTupleValue.of();
+                break;
             default:
                 break;
         }
@@ -398,7 +408,8 @@ protected static int readColumn(String args, int startIndex, int len, String nam
             fixedLength = false;
             estimatedLength++;
         } else if (args.startsWith(matchedKeyword = KEYWORD_TUPLE, i)
-                || args.startsWith(matchedKeyword = KEYWORD_OBJECT, i)) {
+                || args.startsWith(matchedKeyword = KEYWORD_OBJECT, i)
+                || args.startsWith(matchedKeyword = KEYWORD_VARIANT, i)) {
             int index = args.indexOf('(', i + matchedKeyword.length());
             if (index < i) {
                 throw new IllegalArgumentException(ERROR_MISSING_NESTED_TYPE);
@@ -410,12 +421,22 @@ protected static int readColumn(String args, int startIndex, int len, String nam
                 if (c == ')') {
                     break;
                 } else if (c != ',' && !Character.isWhitespace(c)) {
+                    String columnName = "";
                     i = readColumn(args, i, endIndex, "", nestedColumns);
                 }
             }
             if (nestedColumns.isEmpty()) {
                 throw new IllegalArgumentException("Tuple should have at least one nested column");
             }
+
+            List<ClickHouseDataType> variantDataTypes = new ArrayList<>();
+            if (matchedKeyword.equals(KEYWORD_VARIANT)) {
+                nestedColumns.sort(Comparator.comparing(o -> o.getDataType().name()));
+                nestedColumns.forEach(c -> {
+                    c.columnName = "v." + c.getDataType().name();
+                    variantDataTypes.add(c.dataType);
+                });
+            }
             column = new ClickHouseColumn(ClickHouseDataType.valueOf(matchedKeyword), name,
                     args.substring(startIndex, endIndex + 1), nullable, lowCardinality, null, nestedColumns);
             for (ClickHouseColumn n : nestedColumns) {
@@ -424,6 +445,7 @@ protected static int readColumn(String args, int startIndex, int len, String nam
                     fixedLength = false;
                 }
             }
+            column.classToVariantOrdNumMap = ClickHouseDataType.buildVariantMapping(variantDataTypes);
         }
 
         if (column == null) {
@@ -627,6 +649,10 @@ public boolean isAggregateFunction() {
 
     }
 
+    public int getVariantOrdNum(Object value) {
+        return classToVariantOrdNumMap.getOrDefault(value.getClass(), -1);
+    }
+
     public boolean isArray() {
         return dataType == ClickHouseDataType.Array;
     }

clickhouse-data/src/main/java/com/clickhouse/data/ClickHouseDataType.java

@@ -8,12 +8,16 @@
 import java.time.LocalDateTime;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
+import java.util.EnumSet;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.UUID;
 
@@ -101,7 +105,77 @@ public enum ClickHouseDataType {
     Nothing(Object.class, false, true, false, 0, 0, 0, 0, 0, true),
     SimpleAggregateFunction(String.class, true, true, false, 0, 0, 0, 0, 0, false),
     // implementation-defined intermediate state
-    AggregateFunction(String.class, true, true, false, 0, 0, 0, 0, 0, true);
+    AggregateFunction(String.class, true, true, false, 0, 0, 0, 0, 0, true),
+    Variant(List.class, true, true, false, 0, 0, 0, 0, 0, true),
+
+    ;
+
+    public static final List<ClickHouseDataType> ORDERED_BY_RANGE_INT_TYPES =
+            Collections.unmodifiableList(Arrays.asList(
+                    Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64, Int128, UInt128, Int256, UInt256
+            ));
+
+    public static final List<ClickHouseDataType> ORDERED_BY_RANGE_DECIMAL_TYPES =
+            Collections.unmodifiableList(Arrays.asList(
+                    Float32, Float64, Decimal32, Decimal64, Decimal128, Decimal256, Decimal
+            ));
+
+    public static Map<Class<?>, Integer> buildVariantMapping(List<ClickHouseDataType> variantDataTypes) {
+        Map<Class<?>, Integer> variantMapping = new HashMap<>();
+
+        TreeMap<ClickHouseDataType, Integer> intTypesMappings = new TreeMap<>(Comparator.comparingInt(ORDERED_BY_RANGE_INT_TYPES::indexOf));
+        TreeMap<ClickHouseDataType, Integer> decTypesMappings = new TreeMap<>(Comparator.comparingInt(ORDERED_BY_RANGE_DECIMAL_TYPES::indexOf));
+
+        for (int ordNum = 0; ordNum < variantDataTypes.size(); ordNum++) {
+            ClickHouseDataType dataType = variantDataTypes.get(ordNum);
+            Set<Class<?>> classSet = DATA_TYPE_TO_CLASS.get(dataType);
+
+            final int finalOrdNum = ordNum;
+            if (classSet != null) {
+                if (ORDERED_BY_RANGE_INT_TYPES.contains(dataType)) {
+                    intTypesMappings.put(dataType, ordNum);
+                } else if (ORDERED_BY_RANGE_DECIMAL_TYPES.contains(dataType)) {
+                    decTypesMappings.put(dataType, ordNum);
+                } else {
+                    classSet.forEach(c -> variantMapping.put(c, finalOrdNum));
+                }
+            }
+        }
+
+        // add numbers mappings
+        for (java.util.Map.Entry<ClickHouseDataType, Integer> entry : intTypesMappings.entrySet()) {
+            DATA_TYPE_TO_CLASS.get(entry.getKey()).forEach(c -> variantMapping.put(c, entry.getValue()));
+        }
+        for (java.util.Map.Entry<ClickHouseDataType, Integer> entry : decTypesMappings.entrySet()) {
+            DATA_TYPE_TO_CLASS.get(entry.getKey()).forEach(c -> variantMapping.put(c, entry.getValue()));
+        }
+
+
+        return variantMapping;
+    }
+
+    static final Map<ClickHouseDataType, Set<Class<?>>> DATA_TYPE_TO_CLASS = dataTypeClassMap();
+    static Map<ClickHouseDataType, Set<Class<?>>> dataTypeClassMap() {
+        Map<ClickHouseDataType, Set<Class<?>>> map = new HashMap<>();
+
+        map.put(UInt256, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class, BigInteger.class))));
+        map.put(Int256, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class, BigInteger.class))));
+        map.put(UInt128, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class, BigInteger.class))));
+        map.put(Int128, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class, BigInteger.class))));
+        map.put(UInt64, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class))));
+        map.put(Int64, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class, long.class, Long.class))));
+        map.put(UInt32, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class ))));
+        map.put(Int32, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class, int.class, Integer.class))));
+        map.put(UInt16, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class))));
+        map.put(Int16, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class, short.class, Short.class))));
+        map.put(UInt8, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class))));
+        map.put(Int8, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(byte.class, Byte.class))));
+
+        map.put(String, Collections.unmodifiableSet(new HashSet<>(Arrays.asList(String.class))));
+
+        return map;
+    }
+
 
     /**
      * Immutable set(sorted) for all aliases.

clickhouse-data/src/main/java/com/clickhouse/data/format/ClickHouseRowBinaryProcessor.java

@@ -245,6 +245,69 @@ public void serialize(ClickHouseValue value, ClickHouseOutputStream output) thro
         }
     }
 
+    public static class VariantDeserializer extends ClickHouseDeserializer.CompositeDeserializer {
+        private final ClickHouseValue[] values;
+        public VariantDeserializer(ClickHouseDataConfig config, ClickHouseColumn column,
+                ClickHouseDeserializer... deserializers) {
+            super(deserializers);
+
+            List<ClickHouseColumn> nestedCols = column.getNestedColumns();
+            int len = nestedCols.size();
+            if (deserializers.length != len) {
+                throw new IllegalArgumentException(
+                        ClickHouseUtils.format("Expect %d deserializers but got %d", len, deserializers.length));
+            }
+            values = new ClickHouseValue[len];
+            for (int i = 0; i < len; i++) {
+                values[i] = nestedCols.get(i).newValue(config);
+            }
+        }
+
+        @Override
+        public ClickHouseValue deserialize(ClickHouseValue ref, ClickHouseInputStream input) throws IOException {
+            int len = values.length;
+            Object[] tupleValues = new Object[len];
+            int ordTypeNum = BinaryStreamUtils.readInt8(input);
+            for (int i = 0; i < len; i++) {
+                if (ordTypeNum == i) {
+                    tupleValues[i] = deserializers[i].deserialize(values[i], input).asObject();
+                } else {
+                    tupleValues[i] = null;
+                }
+            }
+            return ref.update(tupleValues);
+        }
+    }
+
+    public static class VariantSerializer extends ClickHouseSerializer.CompositeSerializer {
+        private final ClickHouseValue[] values;
+
+        public VariantSerializer(ClickHouseDataConfig config, ClickHouseColumn column,
+                ClickHouseSerializer... serializers) {
+            super(serializers);
+
+            List<ClickHouseColumn> nestedCols = column.getNestedColumns();
+            int len = nestedCols.size();
+            if (serializers.length != len) {
+                throw new IllegalArgumentException(
+                        ClickHouseUtils.format("Expect %d serializers but got %d", len, serializers.length));
+            }
+            values = new ClickHouseValue[len];
+            for (int i = 0; i < len; i++) {
+                values[i] = nestedCols.get(i).newValue(config);
+            }
+        }
+
+        @Override
+        public void serialize(ClickHouseValue value, ClickHouseOutputStream output) throws IOException {
+            List<Object> tupleValues = value.asTuple();
+            // TODO: variant index
+            for (int i = 0, len = serializers.length; i < len; i++) {
+                serializers[i].serialize(values[i].update(tupleValues.get(i)), output);
+            }
+        }
+    }
+
     @Override
     protected void readAndFill(ClickHouseRecord r) throws IOException {
         ClickHouseInputStream in = input;
@@ -512,6 +575,10 @@ public ClickHouseDeserializer getDeserializer(ClickHouseDataConfig config, Click
                 }
                 deserializer = new BitmapSerDe(config, column)::deserialize;
                 break;
+            case Variant:
+                deserializer = new VariantDeserializer(config, column,
+                        getDeserializers(config, column.getNestedColumns()));
+                break;
             default:
                 throw new IllegalArgumentException("Unsupported column:" + column.toString());
         }
@@ -668,6 +735,9 @@ public ClickHouseSerializer getSerializer(ClickHouseDataConfig config, ClickHous
                 }
                 serializer = new BitmapSerDe(config, column)::serialize;
                 break;
+            case Variant:
+                serializer = new VariantSerializer(config, column, getSerializers(config, column.getNestedColumns()));
+                break;
             default:
                 throw new IllegalArgumentException("Unsupported column:" + column.toString());
         }

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseStatementTest.java

@@ -1521,4 +1521,32 @@ public void testMaxResultsRows() throws SQLException {
                     "Unexpected exception: " + e.getMessage());
         }
     }
+
+    @Test(groups = "integration")
+    public void testVariantDataType() throws SQLException {
+        String table = "test_variant_type_01";
+        Properties props = new Properties();
+        props.setProperty("custom_settings", "allow_experimental_variant_type=1");
+        props.setProperty(ClickHouseClientOption.COMPRESS.getKey(), "false");
+        try (ClickHouseConnection conn = newConnection(props);
+             ClickHouseStatement s = conn.createStatement()) {
+
+            s.execute("DROP TABLE IF EXISTS " + table);
+            s.execute("CREATE TABLE " + table +" ( id Variant(UInt32, String, UUID), name String) Engine = MergeTree ORDER BY ()");
+
+            s.execute("insert into " + table + " values ( 1, 'just number' )");
+            s.execute("insert into " + table + " values  ( 'i-am-id-01', 'ID as string' ) ");
+            s.execute("insert into " + table + " values  ( generateUUIDv4(), 'ID as UUID' ) ");
+
+            try (ResultSet rs = s.executeQuery("SELECT * FROM " + table)) {
+                while (rs.next()) {
+                    Object variantValue = rs.getObject(1);
+                    Object name = rs.getString(2);
+                    Object variantSubColumn = rs.getObject("v.String");
+                    System.out.println("-> " + name + " : " + variantValue);
+                    System.out.println("sub: " + variantSubColumn);
+                }
+            }
+        }
+    }
 }

client-v2/pom.xml

@@ -135,6 +135,12 @@
             <scope>test</scope>
         </dependency>
 
+        <dependency>
+            <groupId>org.projectlombok</groupId>
+            <artifactId>lombok</artifactId>
+            <version>1.18.36</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

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

@@ -251,6 +251,7 @@ protected void setSchema(TableSchema schema) {
                 case String:
                 case Enum8:
                 case Enum16:
+                case Variant:
                     this.convertions[i] = NumberConverter.NUMBER_CONVERTERS;
                     break;
                 default:

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

@@ -214,6 +214,8 @@ public <T> T readValue(ClickHouseColumn column, Class<?> typeHint) throws IOExce
                     return (T) readValue(column.getNestedColumns().get(0));
                 case AggregateFunction:
                     return (T) readBitmap( column);
+                case Variant:
+                    return (T) readVariant(column);
                 default:
                     throw new IllegalArgumentException("Unsupported data type: " + column.getDataType());
             }
@@ -675,6 +677,11 @@ public Object[] readTuple(ClickHouseColumn column) throws IOException {
         return tuple;
     }
 
+    public Object readVariant(ClickHouseColumn column) throws IOException {
+        int ordNum = readByte();
+        return readValue(column.getNestedColumns().get(ordNum));
+    }
+
     /**
      * Reads a GEO point as an array of two doubles what represents coordinates (X, Y).
      * @return X, Y coordinates