Added Data Source Hints (awslabs#2148)

Trianz-Akshay · Aug 15, 2024 · 27ed586 · 27ed586
1 parent 6899f8b
commit 27ed586
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 0 deletions.
diff --git a/...com/amazonaws/athena/connector/lambda/metadata/optimizations/DataSourceOptimizations.java b/...com/amazonaws/athena/connector/lambda/metadata/optimizations/DataSourceOptimizations.java
@@ -21,6 +21,7 @@
 
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.ComplexExpressionPushdownSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.FilterPushdownSubType;
+import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.HintsSubtype;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.LimitPushdownSubType;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.PushdownSubTypes;
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.TopNPushdownSubType;
@@ -72,6 +73,18 @@ public Map.Entry<String, List<OptimizationSubType>> withSupportedSubTypes(Pushdo
             }
             return new SimpleImmutableEntry<String, List<OptimizationSubType>>(SUPPORTS_COMPLEX_EXPRESSION_PUSHDOWN.getOptimization(), Arrays.stream(subTypesList).map(pushdownSubTypes -> new OptimizationSubType(pushdownSubTypes.getSubType(), pushdownSubTypes.getProperties())).collect(Collectors.toList()));
         }
+    },
+    //Allows Connector to provide some hints to be used by the engine
+    //Currently this only supports data source default collate to match of Athena
+    DATA_SOURCE_HINTS("data_source_hints")
+    {
+        public Map.Entry<String, List<OptimizationSubType>> withSupportedSubTypes(PushdownSubTypes... subTypesList)
+        {
+            if (!Arrays.stream(subTypesList).allMatch(pushdownSubTypes -> pushdownSubTypes instanceof HintsSubtype)) {
+                throw new IllegalArgumentException("Data Source Hints  must contain valid data source hint subtypes.");
+            }
+            return new SimpleImmutableEntry<>(DATA_SOURCE_HINTS.getOptimization(), Arrays.stream(subTypesList).map(pushdownSubTypes -> new OptimizationSubType(pushdownSubTypes.getSubType(), pushdownSubTypes.getProperties())).collect(Collectors.toList()));
+        }
     };
 
     private final String optimization;

diff --git a/...a/com/amazonaws/athena/connector/lambda/metadata/optimizations/pushdown/HintsSubtype.java b/...a/com/amazonaws/athena/connector/lambda/metadata/optimizations/pushdown/HintsSubtype.java
@@ -0,0 +1,40 @@
+/*-
+ * #%L
+ * Amazon Athena Query Federation SDK
+ * %%
+ * Copyright (C) 2019 - 2022 Amazon Web Services
+ * %%
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * #L%
+ */
+package com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown;
+
+public enum HintsSubtype
+    implements PushdownSubTypes
+{
+    //This indicates that the connector does not support the same default collate setting of Athena
+    NON_DEFAULT_COLLATE("non_default_collate");
+
+    private String subType;
+
+    @Override
+    public String getSubType()
+    {
+        return subType;
+    }
+
+    HintsSubtype(String subType)
+    {
+        this.subType = subType;
+    }
+}