Add support for get partition names by filter

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/CachingHiveMetastore.java

@@ -17,6 +17,7 @@
 import com.facebook.presto.hive.HiveType;
 import com.facebook.presto.hive.MetastoreClientConfig;
 import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.predicate.Domain;
 import com.facebook.presto.spi.security.PrestoPrincipal;
 import com.facebook.presto.spi.security.RoleGrant;
 import com.facebook.presto.spi.statistics.ColumnStatisticType;
@@ -80,7 +81,7 @@ public class CachingHiveMetastore
     private final LoadingCache<HivePartitionName, PartitionStatistics> partitionStatisticsCache;
     private final LoadingCache<String, Optional<List<String>>> viewNamesCache;
     private final LoadingCache<HivePartitionName, Optional<Partition>> partitionCache;
-    private final LoadingCache<PartitionFilter, Optional<List<String>>> partitionFilterCache;
+    private final LoadingCache<PartitionFilter, List<String>> partitionFilterCache;
     private final LoadingCache<HiveTableName, Optional<List<String>>> partitionNamesCache;
     private final LoadingCache<UserTableKey, Set<HivePrivilegeInfo>> tablePrivilegesCache;
     private final LoadingCache<String, Set<String>> rolesCache;
@@ -167,7 +168,7 @@ public Map<HivePartitionName, PartitionStatistics> loadAll(Iterable<? extends Hi
                 .build(asyncReloading(CacheLoader.from(this::loadPartitionNames), executor));
 
         partitionFilterCache = newCacheBuilder(expiresAfterWriteMillis, refreshMills, maximumSize)
-                .build(asyncReloading(CacheLoader.from(this::loadPartitionNamesByParts), executor));
+                .build(asyncReloading(CacheLoader.from(this::loadPartitionNamesByFilter), executor));
 
         partitionCache = newCacheBuilder(expiresAfterWriteMillis, refreshMills, maximumSize)
                 .build(asyncReloading(new CacheLoader<HivePartitionName, Optional<Partition>>()
@@ -524,17 +525,22 @@ private Optional<List<String>> loadPartitionNames(HiveTableName hiveTableName)
     }
 
     @Override
-    public Optional<List<String>> getPartitionNamesByParts(String databaseName, String tableName, List<String> parts)
+    public List<String> getPartitionNamesByFilter(
+            String databaseName,
+            String tableName,
+            Map<Column, Domain> partitionPredicates)
     {
-        return get(partitionFilterCache, partitionFilter(databaseName, tableName, parts));
+        return get(
+                partitionFilterCache,
+                partitionFilter(databaseName, tableName, partitionPredicates));
     }
 
-    private Optional<List<String>> loadPartitionNamesByParts(PartitionFilter partitionFilter)
+    private List<String> loadPartitionNamesByFilter(PartitionFilter partitionFilter)
     {
-        return delegate.getPartitionNamesByParts(
+        return delegate.getPartitionNamesByFilter(
                 partitionFilter.getHiveTableName().getDatabaseName(),
                 partitionFilter.getHiveTableName().getTableName(),
-                partitionFilter.getParts());
+                partitionFilter.getPartitionPredicates());
     }
 
     @Override

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/ExtendedHiveMetastore.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.hive.metastore;
 
 import com.facebook.presto.hive.HiveType;
+import com.facebook.presto.spi.predicate.Domain;
 import com.facebook.presto.spi.security.PrestoPrincipal;
 import com.facebook.presto.spi.security.RoleGrant;
 import com.facebook.presto.spi.statistics.ColumnStatisticType;
@@ -76,7 +77,10 @@ public interface ExtendedHiveMetastore
 
     Optional<List<String>> getPartitionNames(String databaseName, String tableName);
 
-    Optional<List<String>> getPartitionNamesByParts(String databaseName, String tableName, List<String> parts);
+    List<String> getPartitionNamesByFilter(
+            String databaseName,
+            String tableName,
+            Map<Column, Domain> partitionPredicates);
 
     Map<String, Optional<Partition>> getPartitionsByNames(String databaseName, String tableName, List<String> partitionNames);
 

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/MetastoreUtil.java

@@ -24,6 +24,7 @@
 import com.facebook.presto.spi.StandardErrorCode;
 import com.facebook.presto.spi.TableNotFoundException;
 import com.facebook.presto.spi.block.Block;
+import com.facebook.presto.spi.predicate.Domain;
 import com.facebook.presto.spi.type.BigintType;
 import com.facebook.presto.spi.type.BooleanType;
 import com.facebook.presto.spi.type.CharType;
@@ -42,6 +43,7 @@
 import com.facebook.presto.spi.type.VarcharType;
 import com.google.common.base.CharMatcher;
 import com.google.common.collect.ImmutableList;
+import io.airlift.slice.Slice;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.fs.permission.FsPermission;
@@ -50,6 +52,8 @@
 import org.apache.hadoop.hive.metastore.ProtectMode;
 import org.apache.hadoop.io.Text;
 import org.joda.time.DateTimeZone;
+import org.joda.time.format.DateTimeFormatter;
+import org.joda.time.format.ISODateTimeFormat;
 
 import java.io.IOException;
 import java.math.BigInteger;
@@ -60,6 +64,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Properties;
@@ -68,6 +73,7 @@
 import static com.facebook.presto.hive.MetastoreErrorCode.HIVE_FILESYSTEM_ERROR;
 import static com.facebook.presto.hive.MetastoreErrorCode.HIVE_INVALID_PARTITION_VALUE;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED;
+import static com.facebook.presto.spi.type.Chars.padSpaces;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static com.google.common.base.Strings.padEnd;
@@ -102,6 +108,7 @@ public class MetastoreUtil
     public static final String HIVE_DEFAULT_DYNAMIC_PARTITION = "__HIVE_DEFAULT_PARTITION__";
     @SuppressWarnings("OctalInteger")
     public static final FsPermission ALL_PERMISSIONS = new FsPermission((short) 0777);
+    private static final String PARTITION_VALUE_WILDCARD = "";
 
     private MetastoreUtil()
     {
@@ -186,7 +193,7 @@ public static Properties getHiveSchema(
             schema.setProperty(BUCKET_COUNT, "0");
         }
 
-        for (Map.Entry<String, String> param : storage.getSerdeParameters().entrySet()) {
+        for (Entry<String, String> param : storage.getSerdeParameters().entrySet()) {
             schema.setProperty(param.getKey(), (param.getValue() != null) ? param.getValue() : "");
         }
         schema.setProperty(SERIALIZATION_LIB, storage.getStorageFormat().getSerDe());
@@ -234,7 +241,7 @@ public static Properties getHiveSchema(
         }
 
         if (tableParameters != null) {
-            for (Map.Entry<String, String> entry : tableParameters.entrySet()) {
+            for (Entry<String, String> entry : tableParameters.entrySet()) {
                 // add non-null parameters to the schema
                 if (entry.getValue() != null) {
                     schema.setProperty(entry.getKey(), entry.getValue());
@@ -248,10 +255,10 @@ public static Properties getHiveSchema(
     public static Properties getHiveSchema(Map<String, String> serdeParameters, Map<String, String> tableParameters)
     {
         Properties schema = new Properties();
-        for (Map.Entry<String, String> param : serdeParameters.entrySet()) {
+        for (Entry<String, String> param : serdeParameters.entrySet()) {
             schema.setProperty(param.getKey(), (param.getValue() != null) ? param.getValue() : "");
         }
-        for (Map.Entry<String, String> entry : tableParameters.entrySet()) {
+        for (Entry<String, String> entry : tableParameters.entrySet()) {
             // add non-null parameters to the schema
             if (entry.getValue() != null) {
                 schema.setProperty(entry.getKey(), entry.getValue());
@@ -618,4 +625,72 @@ private static String getRenameErrorMessage(Path source, Path target)
     {
         return format("Error moving data files from %s to final location %s", source, target);
     }
+
+    public static List<String> convertPredicateToParts(Map<Column, Domain> partitionPredicates)
+    {
+        List<String> filter = new ArrayList<>();
+
+        for (Entry<Column, Domain> partitionPredicate : partitionPredicates.entrySet()) {
+            Domain domain = partitionPredicate.getValue();
+            if (!domain.isAll()) {
+                if (domain != null && domain.isNullableSingleValue()) {
+                    Object value = domain.getNullableSingleValue();
+                    Type type = domain.getType();
+                    filter.add(convertRawValueToString(value, type));
+                }
+                else {
+                    filter.add(PARTITION_VALUE_WILDCARD);
+                }
+            }
+            else {
+                filter.add(PARTITION_VALUE_WILDCARD);
+            }
+        }
+
+        return filter;
+    }
+
+    public static String convertRawValueToString(Object value, Type type)
+    {
+        String val;
+        if (value == null) {
+            val = HIVE_DEFAULT_DYNAMIC_PARTITION;
+        }
+        else if (type instanceof CharType) {
+            Slice slice = (Slice) value;
+            val = padSpaces(slice, type).toStringUtf8();
+        }
+        else if (type instanceof VarcharType) {
+            Slice slice = (Slice) value;
+            val = slice.toStringUtf8();
+        }
+        else if (type instanceof DecimalType && !((DecimalType) type).isShort()) {
+            Slice slice = (Slice) value;
+            val = Decimals.toString(slice, ((DecimalType) type).getScale());
+        }
+        else if (type instanceof DecimalType && ((DecimalType) type).isShort()) {
+            val = Decimals.toString((long) value, ((DecimalType) type).getScale());
+        }
+        else if (type instanceof DateType) {
+            DateTimeFormatter dateTimeFormatter = ISODateTimeFormat.date().withZoneUTC();
+            val = dateTimeFormatter.print(TimeUnit.DAYS.toMillis((long) value));
+        }
+        else if (type instanceof TimestampType) {
+            // we don't have time zone info, so just add a wildcard
+            val = PARTITION_VALUE_WILDCARD;
+        }
+        else if (type instanceof TinyintType
+                || type instanceof SmallintType
+                || type instanceof IntegerType
+                || type instanceof BigintType
+                || type instanceof DoubleType
+                || type instanceof RealType
+                || type instanceof BooleanType) {
+            val = value.toString();
+        }
+        else {
+            throw new PrestoException(NOT_SUPPORTED, format("Unsupported partition key type: %s", type.getDisplayName()));
+        }
+        return val;
+    }
 }

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/PartitionFilter.java

@@ -13,13 +13,11 @@
  */
 package com.facebook.presto.hive.metastore;
 
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.google.common.collect.ImmutableList;
+import com.facebook.presto.spi.predicate.Domain;
 
 import javax.annotation.concurrent.Immutable;
 
-import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 import static com.facebook.presto.hive.metastore.HiveTableName.hiveTableName;
@@ -30,38 +28,36 @@
 public class PartitionFilter
 {
     private final HiveTableName hiveTableName;
-    private final List<String> parts;
 
-    @JsonCreator
-    public PartitionFilter(@JsonProperty("hiveTableName") HiveTableName hiveTableName, @JsonProperty("parts") List<String> parts)
+    private final Map<Column, Domain> partitionPredicates;
+
+    public PartitionFilter(HiveTableName hiveTableName, Map<Column, Domain> partitionPredicates)
     {
         this.hiveTableName = requireNonNull(hiveTableName, "hiveTableName is null");
-        this.parts = ImmutableList.copyOf(requireNonNull(parts, "parts is null"));
+        this.partitionPredicates = requireNonNull(partitionPredicates, "effectivePredicate is null");
     }
 
-    public static PartitionFilter partitionFilter(String databaseName, String tableName, List<String> parts)
+    public static PartitionFilter partitionFilter(String databaseName, String tableName, Map<Column, Domain> effectivePredicate)
     {
-        return new PartitionFilter(hiveTableName(databaseName, tableName), parts);
+        return new PartitionFilter(hiveTableName(databaseName, tableName), effectivePredicate);
     }
 
-    @JsonProperty
     public HiveTableName getHiveTableName()
     {
         return hiveTableName;
     }
 
-    @JsonProperty
-    public List<String> getParts()
+    public Map<Column, Domain> getPartitionPredicates()
     {
-        return parts;
+        return partitionPredicates;
     }
 
     @Override
     public String toString()
     {
         return toStringHelper(this)
                 .add("hiveTableName", hiveTableName)
-                .add("parts", parts)
+                .add("partitionPredicates", partitionPredicates)
                 .toString();
     }
 
@@ -77,12 +73,12 @@ public boolean equals(Object o)
 
         PartitionFilter other = (PartitionFilter) o;
         return Objects.equals(hiveTableName, other.hiveTableName) &&
-                Objects.equals(parts, other.parts);
+                Objects.equals(partitionPredicates, other.partitionPredicates);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(hiveTableName, parts);
+        return Objects.hash(hiveTableName, partitionPredicates);
     }
 }