From 9ae10f0fbc5e2daa2108a75a4f0dd8b06edbed35 Mon Sep 17 00:00:00 2001
From: Yuya Ebihara <ebyhry@gmail.com>
Date: Mon, 7 Sep 2020 05:08:22 -0700
Subject: [PATCH] Skip unsupported types when pushing down in Cassandra

---
 .../cassandra/CassandraPartitionManager.java       |  4 +++-
 .../facebook/presto/cassandra/CassandraType.java   |  7 +++----
 .../presto/cassandra/CassandraTestingUtils.java    | 14 +++++++-------
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraPartitionManager.java b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraPartitionManager.java
index fd03ed008673..7c68969c8f0f 100644
--- a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraPartitionManager.java
+++ b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraPartitionManager.java
@@ -152,7 +152,9 @@ private static List<Set<Object>> getPartitionKeysList(CassandraTable table, Tupl
                             Object value = range.getSingleValue();
 
                             CassandraType valueType = columnHandle.getCassandraType();
-                            columnValues.add(valueType.validatePartitionKey(value));
+                            if (valueType.isSupportedPartitionKey()) {
+                                columnValues.add(value);
+                            }
                         }
                         return columnValues.build();
                     },
diff --git a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraType.java b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraType.java
index e6128573f73d..6b615cb5b261 100644
--- a/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraType.java
+++ b/presto-cassandra/src/main/java/com/facebook/presto/cassandra/CassandraType.java
@@ -445,7 +445,7 @@ public Object getJavaValue(Object nativeValue)
         }
     }
 
-    public Object validatePartitionKey(Object value)
+    public boolean isSupportedPartitionKey()
     {
         switch (this) {
             case ASCII:
@@ -461,7 +461,7 @@ public Object validatePartitionKey(Object value)
             case TIMESTAMP:
             case UUID:
             case TIMEUUID:
-                return value;
+                return true;
             case COUNTER:
             case BLOB:
             case CUSTOM:
@@ -470,8 +470,7 @@ public Object validatePartitionKey(Object value)
             case LIST:
             case MAP:
             default:
-                // todo should we just skip partition pruning instead of throwing an exception?
-                throw new PrestoException(NOT_SUPPORTED, "Unsupport partition key type: " + this);
+                return false;
         }
     }
 
diff --git a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/CassandraTestingUtils.java b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/CassandraTestingUtils.java
index fff1018791b0..791f72c53301 100644
--- a/presto-cassandra/src/test/java/com/facebook/presto/cassandra/CassandraTestingUtils.java
+++ b/presto-cassandra/src/test/java/com/facebook/presto/cassandra/CassandraTestingUtils.java
@@ -197,23 +197,23 @@ public static void createTableAllTypesPartitionKey(CassandraSession session, Sch
                 "   typeinteger, " +
                 "   typelong, " +
                 // TODO: NOT YET SUPPORTED AS A PARTITION KEY
-                // "   typebytes, " +
+                "   typebytes, " +
                 "   typetimestamp, " +
                 "   typeansi, " +
                 "   typeboolean, " +
                 // TODO: PRECISION LOST. IMPLEMENT IT AS STRING
-                //  "   typedecimal, " +
+                "   typedecimal, " +
                 "   typedouble, " +
                 "   typefloat, " +
                 "   typeinet, " +
                 "   typevarchar, " +
                 // TODO: NOT YET SUPPORTED AS A PARTITION KEY
-                // "   typevarint, " +
-                "   typetimeuuid " +
+                "   typevarint, " +
+                "   typetimeuuid, " +
                 // TODO: NOT YET SUPPORTED AS A PARTITION KEY
-                // "   typelist, " +
-                // "   typemap, " +
-                // "   typeset" +
+                "   typelist, " +
+                "   typemap, " +
+                "   typeset" +
                 " ))" +
                 ")");