Move partition info to CassandraType.Kind

trinodb · Sep 27, 2023 · 609ac97 · 609ac97
1 parent 4752df4
commit 609ac97
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 62 deletions.
diff --git a/...in/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraPartitionManager.java b/...in/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraPartitionManager.java
@@ -160,7 +160,7 @@ private List<Set<Object>> getPartitionKeysList(CassandraTable table, TupleDomain
                             Object value = range.getSingleValue();
 
                             CassandraType valueType = columnHandle.getCassandraType();
-                            if (cassandraTypeManager.isSupportedPartitionKey(valueType.getKind())) {
+                            if (valueType.getKind().isSupportedPartitionKey()) {
                                 columnValues.add(value);
                             }
                         }

diff --git a/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraType.java b/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraType.java
@@ -28,32 +28,45 @@ public class CassandraType
 {
     public enum Kind
     {
-        BOOLEAN,
-        TINYINT,
-        SMALLINT,
-        INT,
-        BIGINT,
-        FLOAT,
-        DOUBLE,
-        DECIMAL,
-        DATE,
-        TIME,
-        TIMESTAMP,
-        ASCII,
-        TEXT,
-        VARCHAR,
-        BLOB,
-        UUID,
-        TIMEUUID,
-        COUNTER,
-        VARINT,
-        INET,
-        CUSTOM,
-        LIST,
-        SET,
-        MAP,
-        TUPLE,
-        UDT,
+        BOOLEAN(true),
+        TINYINT(true),
+        SMALLINT(true),
+        INT(true),
+        BIGINT(true),
+        FLOAT(true),
+        DOUBLE(true),
+        DECIMAL(true),
+        DATE(true),
+        TIME(true),
+        TIMESTAMP(true),
+        ASCII(true),
+        TEXT(true),
+        VARCHAR(true),
+        BLOB(false),
+        UUID(true),
+        TIMEUUID(true),
+        COUNTER(false),
+        VARINT(false),
+        INET(true),
+        CUSTOM(false),
+        LIST(false),
+        SET(false),
+        MAP(false),
+        TUPLE(false),
+        UDT(false),
+        /**/;
+
+        private final boolean supportedPartitionKey;
+
+        Kind(boolean supportedPartitionKey)
+        {
+            this.supportedPartitionKey = supportedPartitionKey;
+        }
+
+        public boolean isSupportedPartitionKey()
+        {
+            return supportedPartitionKey;
+        }
     }
 
     private final Kind kind;

diff --git a/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java b/plugin/trino-cassandra/src/main/java/io/trino/plugin/cassandra/CassandraTypeManager.java
@@ -487,41 +487,6 @@ public Object getJavaValue(CassandraType.Kind kind, Object trinoNativeValue)
         throw new IllegalStateException("Back conversion not implemented for " + this);
     }
 
-    public boolean isSupportedPartitionKey(CassandraType.Kind kind)
-    {
-        switch (kind) {
-            case ASCII:
-            case TEXT:
-            case VARCHAR:
-            case BIGINT:
-            case BOOLEAN:
-            case DOUBLE:
-            case INET:
-            case INT:
-            case TINYINT:
-            case SMALLINT:
-            case FLOAT:
-            case DECIMAL:
-            case DATE:
-            case TIME:
-            case TIMESTAMP:
-            case UUID:
-            case TIMEUUID:
-                return true;
-            case COUNTER:
-            case BLOB:
-            case CUSTOM:
-            case VARINT:
-            case SET:
-            case LIST:
-            case MAP:
-            case TUPLE:
-            case UDT:
-            default:
-                return false;
-        }
-    }
-
     public boolean isFullySupported(DataType dataType)
     {
         if (toCassandraType(dataType).isEmpty()) {