feat: Support metadata projection

presto-clp/pom.xml

@@ -18,6 +18,19 @@
     </properties>
 
     <dependencies>
+        <dependency>
+            <groupId>org.msgpack</groupId>
+            <artifactId>msgpack-core</artifactId>
+            <version>0.9.8</version>
+        </dependency>
+
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>4.6.1</version>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.mysql</groupId>
             <artifactId>mysql-connector-j</artifactId>

presto-clp/src/main/java/com/facebook/presto/plugin/clp/ClpErrorCode.java

@@ -29,9 +29,11 @@ public enum ClpErrorCode
     CLP_UNSUPPORTED_TYPE(3, EXTERNAL),
     CLP_UNSUPPORTED_CONFIG_OPTION(4, EXTERNAL),
     CLP_UNSUPPORTED_TABLE_SCHEMA_YAML(5, EXTERNAL),
+    CLP_UNSUPPORTED_METADATA_PROJECTION(6, EXTERNAL),
 
     CLP_SPLIT_METADATA_CONFIG_NOT_FOUND(10, USER_ERROR),
-    CLP_MANDATORY_COLUMN_NOT_IN_FILTER(11, USER_ERROR);
+    CLP_MANDATORY_COLUMN_NOT_IN_FILTER(11, USER_ERROR),
+    CLP_SPLIT_METADATA_TYPE_MISMATCH_METADATA_DATABASE_TYPE(12, USER_ERROR);
 
     private final ErrorCode errorCode;
 

presto-clp/src/main/java/com/facebook/presto/plugin/clp/ClpMetadata.java

@@ -119,7 +119,7 @@ public ConnectorTableLayoutResult getTableLayoutForConstraint(
             Optional<Set<ColumnHandle>> desiredColumns)
     {
         ClpTableHandle tableHandle = (ClpTableHandle) table;
-        ConnectorTableLayout layout = new ConnectorTableLayout(new ClpTableLayoutHandle(tableHandle, Optional.empty(), Optional.empty()));
+        ConnectorTableLayout layout = new ConnectorTableLayout(new ClpTableLayoutHandle(tableHandle, Optional.empty(), null));
         return new ConnectorTableLayoutResult(layout, constraint.getSummary());
     }
 

presto-clp/src/main/java/com/facebook/presto/plugin/clp/ClpSplit.java

@@ -21,7 +21,12 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import org.msgpack.core.MessageBufferPacker;
+import org.msgpack.core.MessagePack;
 
+import java.io.IOException;
+import java.util.Base64;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -36,16 +41,56 @@ public class ClpSplit
     private final String path;
     private final SplitType type;
     private final Optional<String> kqlQuery;
+    private final String metadataProjectionNameValueEncoded;
 
+    /**
+     * Invoked by Jackson; serializes a ClpSplit to JSON format
+     *
+     * @param path the path to the split
+     * @param type the split type
+     * @param kqlQuery optional KQL query pushed down to CLP-S
+     * @param metadataProjectionNameValueEncoded Base64-encoded MessagePack representation of metadata projection in
+     *                                           column name and value pairs
+     */
     @JsonCreator
     public ClpSplit(
             @JsonProperty("path") String path,
             @JsonProperty("type") SplitType type,
-            @JsonProperty("kqlQuery") Optional<String> kqlQuery)
+            @JsonProperty("kqlQuery") Optional<String> kqlQuery,
+            @JsonProperty("metadataProjectionNameValue") String metadataProjectionNameValueEncoded)
     {
         this.path = requireNonNull(path, "Split path is null");
         this.type = requireNonNull(type, "Split type is null");
         this.kqlQuery = kqlQuery;
+        this.metadataProjectionNameValueEncoded = metadataProjectionNameValueEncoded != null ? metadataProjectionNameValueEncoded : "";
+    }
+
+    /**
+     * Creates a ClpSplit for internal use
+     *
+     * @param path the path to the split
+     * @param type the split type
+     * @param kqlQuery optional KQL query pushed down to CLP-S
+     * @param metadataProjectionNameValue optional map of metadata projection column names to their values
+     * @throws RuntimeException if encoding metadata projection name-value pairs fails
+     */
+    public ClpSplit(
+            String path,
+            SplitType type,
+            Optional<String> kqlQuery,
+            Optional<Map<String, Object>> metadataProjectionNameValue)
+    {
+        this.path = requireNonNull(path, "Split path is null");
+        this.type = requireNonNull(type, "Split type is null");
+        this.kqlQuery = kqlQuery;
+
+        try {
+            this.metadataProjectionNameValueEncoded = encodeProjectionNameValue(
+                    metadataProjectionNameValue.orElse(new HashMap<>()));
+        }
+        catch (IOException e) {
+            throw new RuntimeException("Failed to encode projection name value", e);
+        }
     }
 
     @JsonProperty
@@ -66,6 +111,48 @@ public Optional<String> getKqlQuery()
         return kqlQuery;
     }
 
+    @JsonProperty("metadataProjectionNameValue")
+    public String getMetadataProjectionNameValue()
+    {
+        return metadataProjectionNameValueEncoded;
+    }
+
+    /**
+     * @param metadataColumnNameValue map of metadata column names to their values.
+     * @return Base64-encoded MessagePack representation of the metadataColumnNameValue map
+     * @throws IOException if MessagePack serialization fails
+     * @throws IllegalArgumentException if a value type is not String, Long, or Double
+     */
+    private String encodeProjectionNameValue(Map<String, Object> metadataColumnNameValue) throws IOException
+    {
+        if (metadataColumnNameValue.isEmpty()) {
+            return "";
+        }
+
+        try (MessageBufferPacker packer = MessagePack.newDefaultBufferPacker()) {
+            packer.packMapHeader(metadataColumnNameValue.size());
+            for (Map.Entry<String, Object> entry : metadataColumnNameValue.entrySet()) {
+                packer.packString(entry.getKey());
+
+                Object value = entry.getValue();
+                if (value instanceof String) {
+                    packer.packString((String) value);
+                }
+                else if (value instanceof Long) {
+                    packer.packLong((Long) value);
+                }
+                else if (value instanceof Double) {
+                    packer.packDouble((Double) value);
+                }
+                else {
+                    throw new IllegalArgumentException(
+                            "Unsupported type for column " + entry.getKey() + ": " + value.getClass());
+                }
+            }
+            return Base64.getEncoder().encodeToString(packer.toByteArray());
+        }
+    }
+
     @Override
     public NodeSelectionStrategy getNodeSelectionStrategy()
     {

presto-clp/src/main/java/com/facebook/presto/plugin/clp/ClpTableLayoutHandle.java

@@ -19,8 +19,10 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
+import java.util.HashSet;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 
@@ -29,34 +31,51 @@ public class ClpTableLayoutHandle
 {
     private final ClpTableHandle table;
     private final Optional<String> kqlQuery;
-    private final Optional<RowExpression> metadataExpression;
+    private final RowExpression metadataExpression;
     private final boolean metadataQueryOnly;
     private final Optional<ClpTopNSpec> topN;
 
+    private Optional<Set<String>> splitMetadataColumnNames;
+
     @JsonCreator
     public ClpTableLayoutHandle(
             @JsonProperty("table") ClpTableHandle table,
             @JsonProperty("kqlQuery") Optional<String> kqlQuery,
-            @JsonProperty("metadataExpression") Optional<RowExpression> metadataExpression,
+            @JsonProperty("metadataExpression") RowExpression metadataExpression,
             @JsonProperty("metadataQueryOnly") boolean metadataQueryOnly,
+            @JsonProperty("splitMetadataColumnNames") Optional<Set<String>> splitMetadataColumnNames,
             @JsonProperty("topN") Optional<ClpTopNSpec> topN)
     {
         this.table = table;
         this.kqlQuery = kqlQuery;
         this.metadataExpression = metadataExpression;
         this.metadataQueryOnly = metadataQueryOnly;
+        this.splitMetadataColumnNames = splitMetadataColumnNames;
         this.topN = topN;
     }
 
     public ClpTableLayoutHandle(
             @JsonProperty("table") ClpTableHandle table,
             @JsonProperty("kqlQuery") Optional<String> kqlQuery,
-            @JsonProperty("metadataExpression") Optional<RowExpression> metadataExpression)
+            @JsonProperty("metadataExpression") RowExpression metadataExpression)
     {
         this.table = table;
         this.kqlQuery = kqlQuery;
         this.metadataExpression = metadataExpression;
         this.metadataQueryOnly = false;
+        this.splitMetadataColumnNames = Optional.empty();
+        this.topN = Optional.empty();
+    }
+
+    public ClpTableLayoutHandle(
+            @JsonProperty("table") ClpTableHandle table,
+            @JsonProperty("splitMetadataColumnNames") Optional<Set<String>> splitMetadataColumnNames)
+    {
+        this.table = table;
+        this.kqlQuery = Optional.empty();
+        this.metadataExpression = null;
+        this.metadataQueryOnly = false;
+        this.splitMetadataColumnNames = splitMetadataColumnNames;
         this.topN = Optional.empty();
     }
 
@@ -73,7 +92,7 @@ public Optional<String> getKqlQuery()
     }
 
     @JsonProperty
-    public Optional<RowExpression> getMetadataExpression()
+    public RowExpression getMetadataExpression()
     {
         return metadataExpression;
     }
@@ -84,6 +103,20 @@ public boolean isMetadataQueryOnly()
         return metadataQueryOnly;
     }
 
+    @JsonProperty
+    public Optional<Set<String>> getSplitMetadataColumnNames()
+    {
+        return splitMetadataColumnNames;
+    }
+
+    public Set<String> getOrInitializeSplitMetadataColumnNames()
+    {
+        if (!splitMetadataColumnNames.isPresent()) {
+            splitMetadataColumnNames = Optional.of(new HashSet<>());
+        }
+        return splitMetadataColumnNames.get();
+    }
+
     @JsonProperty
     public Optional<ClpTopNSpec> getTopN()
     {
@@ -103,14 +136,15 @@ public boolean equals(Object o)
         return Objects.equals(table, that.table) &&
                 Objects.equals(kqlQuery, that.kqlQuery) &&
                 Objects.equals(metadataExpression, that.metadataExpression) &&
-                Objects.equals(metadataQueryOnly, that.metadataQueryOnly) &&
+                Objects.equals(splitMetadataColumnNames, that.splitMetadataColumnNames) &&
                 Objects.equals(topN, that.topN);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(table, kqlQuery, metadataExpression, metadataQueryOnly, topN);
+        return Objects.hash(
+                table, kqlQuery, metadataExpression, metadataQueryOnly, topN);
     }
 
     @Override