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/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,52 @@ 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 "";
+        }
+
+        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());
+            }
+        }
+
+        byte[] bytes = packer.toByteArray();
+        packer.close();
+
+        return Base64.getEncoder().encodeToString(bytes);
+    }
+
     @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;
 
@@ -33,18 +35,22 @@ public class ClpTableLayoutHandle
     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("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;
     }
 
@@ -57,6 +63,19 @@ public ClpTableLayoutHandle(
         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 = Optional.empty();
+        this.metadataQueryOnly = false;
+        this.splitMetadataColumnNames = splitMetadataColumnNames;
         this.topN = Optional.empty();
     }
 
@@ -84,6 +103,15 @@ public boolean isMetadataQueryOnly()
         return metadataQueryOnly;
     }
 
+    @JsonProperty
+    public Set<String> getSplitMetadataColumnNames()
+    {
+        if (!splitMetadataColumnNames.isPresent()) {
+            splitMetadataColumnNames = Optional.of(new HashSet<>());
+        }
+        return splitMetadataColumnNames.get();
+    }
+
     @JsonProperty
     public Optional<ClpTopNSpec> getTopN()
     {
@@ -104,13 +132,15 @@ public boolean equals(Object o)
                 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, splitMetadataColumnNames, topN);
     }
 
     @Override
@@ -121,6 +151,7 @@ public String toString()
                 .add("kqlQuery", kqlQuery)
                 .add("metadataExpression", metadataExpression)
                 .add("metadataQueryOnly", metadataQueryOnly)
+                .add("splitMetaColumnNames", splitMetadataColumnNames)
                 .add("topN", topN)
                 .toString();
     }

presto-clp/src/main/java/com/facebook/presto/plugin/clp/optimization/ClpComputePushDown.java

@@ -26,6 +26,7 @@
 import com.facebook.presto.spi.ConnectorPlanRewriter;
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.ConnectorTableLayoutHandle;
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.TableHandle;
 import com.facebook.presto.spi.VariableAllocator;
@@ -47,12 +48,14 @@
 import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Deque;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 
+import static com.facebook.presto.plugin.clp.ClpErrorCode.CLP_UNSUPPORTED_METADATA_PROJECTION;
 import static com.facebook.presto.spi.ConnectorPlanRewriter.rewriteWith;
 import static java.lang.Math.toIntExact;
 import static java.lang.String.format;
@@ -100,7 +103,91 @@ public PlanNode visitFilter(FilterNode node, RewriteContext<Void> context)
                 return node;
             }
 
-            return processFilter(node, (TableScanNode) node.getSource());
+            PlanNode processedNode = processFilter(node, (TableScanNode) node.getSource());
+
+            if (processedNode instanceof TableScanNode) {
+                return context.rewrite(processedNode, null);
+            }
+
+            return processedNode;
+        }
+
+        /**
+         * Rewrites a TableScanNode to attach metadata projection column.
+         *
+         * @param node the original TableScanNode to rewrite.
+         * @param context
+         * @return a new TableScanNode with metadata projection in the layout handle
+         * @throw PrestoException if a metadata column maps to a range-bound column - this is an unsupported feature
+         */
+        @Override
+        public PlanNode visitTableScan(TableScanNode node, RewriteContext<Void> context)
+        {
+            // Retrieve projection column names
+            Set<String> projectionColumns = new HashSet<>();
+            for (VariableReferenceExpression variable : node.getOutputVariables()) {
+                projectionColumns.add(variable.getName());
+            }
+
+            // Retrieve metadata column
+            TableHandle tableHandle = node.getTable();
+            ClpTableHandle clpTableHandle = (ClpTableHandle) tableHandle.getConnectorHandle();
+            SchemaTableName schemaTableName = clpTableHandle.getSchemaTableName();
+            Set<String> metadataColumns = metadataConfig.getMetadataColumns(schemaTableName).keySet();
+
+            // Metadata Projection: intersection between the projection column and metadata column
+            Set<String> metadataProjections = new HashSet<>();
+            for (String columnName : projectionColumns) {
+                if (metadataColumns.contains(columnName)) {
+                    Set<String> metadataColumnsWithRangeBound =
+                            metadataConfig.getMetadataColumnsWithRangeBounds(schemaTableName);
+                    Map<String, String> exposedToOriginalMap =
+                            metadataConfig.getExposedToOriginalMapping(schemaTableName);
+
+                    // Resolve exposed column names to their original names in the metadata database.
+                    // After extracting values from the metadata database, these will be mapped back to exposed names
+                    // for projection.
+                    String originalColumnName = exposedToOriginalMap.get(columnName);
+                    if (metadataColumnsWithRangeBound.contains(originalColumnName)) {
+                        throw new PrestoException(CLP_UNSUPPORTED_METADATA_PROJECTION,
+                                format("Unsupported metadata projection column: %s", columnName));
+                    }
+                    metadataProjections.add(originalColumnName);
+                }
+            }
+
+            if (metadataProjections.isEmpty()) {
+                return node;
+            }
+
+            // TableScan optimization happens late in planning; append to existing layout if present.
+            Optional<ConnectorTableLayoutHandle> layout = tableHandle.getLayout();
+            if (layout.isPresent() && layout.get() instanceof ClpTableLayoutHandle) {
+                ClpTableLayoutHandle cl = (ClpTableLayoutHandle) layout.get();
+                for (String metadataProjection : metadataProjections) {
+                    cl.getSplitMetadataColumnNames().add(metadataProjection);
+                }
+                return node;
+            }
+
+            ClpTableLayoutHandle newLayout = new ClpTableLayoutHandle(
+                    clpTableHandle, Optional.of(metadataProjections));
+
+            // TableScanNode is immutable, we need to copy-on-write
+            return new TableScanNode(
+                    node.getSourceLocation(),
+                    idAllocator.getNextId(),
+                    new TableHandle(
+                            tableHandle.getConnectorId(),
+                            clpTableHandle,
+                            tableHandle.getTransaction(),
+                            Optional.of(newLayout)),
+                    node.getOutputVariables(),
+                    node.getAssignments(),
+                    node.getTableConstraints(),
+                    node.getCurrentConstraint(),
+                    node.getEnforcedConstraint(),
+                    node.getCteMaterializationInfo());
         }
 
         @Override
@@ -191,7 +278,7 @@ public PlanNode visitTopN(TopNNode node, RewriteContext<Void> context)
                 ClpTopNSpec tightened = new ClpTopNSpec(mergedLimit, ex.getOrderings());
                 ClpTableHandle clpHandle = (ClpTableHandle) tableHandle.getConnectorHandle();
                 ClpTableLayoutHandle newLayout =
-                        new ClpTableLayoutHandle(clpHandle, kql, metadataSql, true, Optional.of(tightened));
+                        new ClpTableLayoutHandle(clpHandle, kql, metadataSql, true, Optional.empty(), Optional.of(tightened));
 
                 TableScanNode newScan = new TableScanNode(
                         scan.getSourceLocation(),
@@ -227,7 +314,7 @@ public PlanNode visitTopN(TopNNode node, RewriteContext<Void> context)
             ClpTopNSpec spec = new ClpTopNSpec(node.getCount(), newOrderings);
             ClpTableHandle clpHandle = (ClpTableHandle) tableHandle.getConnectorHandle();
             ClpTableLayoutHandle newLayout =
-                    new ClpTableLayoutHandle(clpHandle, kql, metadataSql, true, Optional.of(spec));
+                    new ClpTableLayoutHandle(clpHandle, kql, metadataSql, true, Optional.empty(), Optional.of(spec));
 
             TableScanNode newScanNode = new TableScanNode(
                     scan.getSourceLocation(),
@@ -283,7 +370,7 @@ private PlanNode processFilter(FilterNode filterNode, TableScanNode tableScanNod
                 kqlQuery.ifPresent(s -> log.debug("KQL query: %s", s));
 
                 ClpTableLayoutHandle layoutHandle = new ClpTableLayoutHandle(
-                        clpTableHandle, kqlQuery, metadataExpression, allInMetadata, Optional.empty());
+                        clpTableHandle, kqlQuery, metadataExpression, allInMetadata, Optional.empty(), Optional.empty());
                 TableHandle newTableHandle = new TableHandle(
                         tableHandle.getConnectorId(),
                         clpTableHandle,