Infer schema from gremlin/sparql query results for Neptune qpt (awslabs#2543)

Co-authored-by: akshay.kachore <[email protected]>

Author: VenkatasivareddyTR

athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/Constants.java

@@ -53,4 +53,7 @@ protected Constants()
    public static final int PREFIX_LEN = PREFIX_KEY.length();
 
    public static final String GREMLIN_QUERY_SUPPORT_TYPE = "valueMap";
+   public static final String RDF_COMPONENT_TYPE = "rdf";
+   public static final String GREMLIN_QUERY_LIMIT = "limit(10)";
+   public static final String SPARQL_QUERY_LIMIT = "limit 10";
 }

athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandler.java

@@ -22,7 +22,6 @@
 import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
 import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
 import com.amazonaws.athena.connector.lambda.data.BlockWriter;
-import com.amazonaws.athena.connector.lambda.data.SchemaBuilder;
 import com.amazonaws.athena.connector.lambda.domain.Split;
 import com.amazonaws.athena.connector.lambda.domain.TableName;
 import com.amazonaws.athena.connector.lambda.domain.spill.SpillLocation;
@@ -42,7 +41,8 @@
 import com.amazonaws.athena.connector.lambda.metadata.optimizations.OptimizationSubType;
 import com.amazonaws.athena.connector.lambda.security.EncryptionKeyFactory;
 import com.amazonaws.athena.connectors.neptune.propertygraph.PropertyGraphHandler;
-import com.amazonaws.athena.connectors.neptune.qpt.NeptuneQueryPassthrough;
+import com.amazonaws.athena.connectors.neptune.qpt.NeptuneGremlinQueryPassthrough;
+import com.amazonaws.athena.connectors.neptune.qpt.NeptuneSparqlQueryPassthrough;
 import com.amazonaws.athena.connectors.neptune.rdf.NeptuneSparqlConnection;
 import com.google.common.collect.ImmutableMap;
 import org.apache.arrow.util.VisibleForTesting;
@@ -67,6 +67,9 @@
 import java.util.NoSuchElementException;
 import java.util.Set;
 
+import static com.amazonaws.athena.connectors.neptune.Constants.GREMLIN_QUERY_LIMIT;
+import static com.amazonaws.athena.connectors.neptune.Constants.RDF_COMPONENT_TYPE;
+import static com.amazonaws.athena.connectors.neptune.Constants.SPARQL_QUERY_LIMIT;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -94,7 +97,8 @@ public class NeptuneMetadataHandler extends GlueMetadataHandler
     private final String glueDBName;
 
     private NeptuneConnection neptuneConnection = null;
-    private final NeptuneQueryPassthrough queryPassthrough = new NeptuneQueryPassthrough();
+    private final NeptuneGremlinQueryPassthrough gremlinQueryPassthrough = new NeptuneGremlinQueryPassthrough();
+    private final NeptuneSparqlQueryPassthrough sparqlQueryPassthrough = new NeptuneSparqlQueryPassthrough();
 
     public NeptuneMetadataHandler(java.util.Map<String, String> configOptions)
     {
@@ -127,16 +131,17 @@ protected NeptuneMetadataHandler(
     public GetDataSourceCapabilitiesResponse doGetDataSourceCapabilities(BlockAllocator allocator, GetDataSourceCapabilitiesRequest request)
     {
         ImmutableMap.Builder<String, List<OptimizationSubType>> capabilities = ImmutableMap.builder();
-        queryPassthrough.addQueryPassthroughCapabilityIfEnabled(capabilities, configOptions);
+        gremlinQueryPassthrough.addQueryPassthroughCapabilityIfEnabled(capabilities, configOptions);
+        sparqlQueryPassthrough.addQueryPassthroughCapabilityIfEnabled(capabilities, configOptions);
 
         return new GetDataSourceCapabilitiesResponse(request.getCatalogName(), capabilities.build());
     }
-    
+
     /**
      * Since the entire Neptune cluster is considered as a single graph database,
      * just return the glue database name provided as a single database (schema)
      * name.
-     * 
+     *
      * @param allocator Tool for creating and managing Apache Arrow Blocks.
      * @param request   Provides details on who made the request and which Athena
      *                  catalog they are querying.
@@ -158,7 +163,7 @@ public ListSchemasResponse doListSchemaNames(BlockAllocator allocator, ListSchem
     /**
      * Used to get the list of tables that this data source contains. In this case,
      * fetch list of tables in the Glue database provided.
-     * 
+     *
      * @param allocator Tool for creating and managing Apache Arrow Blocks.
      * @param request   Provides details on who made the request and which Athena
      *                  catalog and database they are querying.
@@ -191,7 +196,7 @@ public ListTablesResponse doListTables(BlockAllocator allocator, ListTablesReque
     /**
      * Used to get definition (field names, types, descriptions, etc...) of a Table.
      *
-     * @param allocator Tool for creating and managing Apache Arrow Blocks.
+     * @param blockAllocator Tool for creating and managing Apache Arrow Blocks.
      * @param request   Provides details on who made the request and which Athena
      *                  catalog, database, and table they are querying.
      * @return A GetTableResponse which primarily contains: 1. An Apache Arrow
@@ -209,10 +214,10 @@ public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableReques
         Schema tableSchema = null;
         try {
             if (glue != null) {
-                tableSchema = super.doGetTable(blockAllocator, request).getSchema();        
+                tableSchema = super.doGetTable(blockAllocator, request).getSchema();
                 logger.info("doGetTable: Retrieved schema for table[{}] from AWS Glue.", request.getTableName());
             }
-        } 
+        }
         catch (RuntimeException ex) {
             logger.warn("doGetTable: Unable to retrieve table[{}:{}] from AWS Glue.",
                     request.getTableName().getSchemaName(), request.getTableName().getTableName(), ex);
@@ -226,7 +231,7 @@ public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableReques
      */
     @Override
     public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request,
-            QueryStatusChecker queryStatusChecker) throws Exception 
+            QueryStatusChecker queryStatusChecker) throws Exception
     {
         // No implemenation as connector doesn't support partitioning
     }
@@ -255,7 +260,7 @@ public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request
      * the RecordHandler has easy access to it.
      */
     @Override
-    public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest request) 
+    public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest request)
     {
         // Every split must have a unique location if we wish to spill to avoid failures
         SpillLocation spillLocation = makeSpillLocation(request);
@@ -266,7 +271,7 @@ public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsReq
     }
 
     @Override
-    protected Field convertField(String name, String glueType) 
+    protected Field convertField(String name, String glueType)
     {
         return GlueFieldLexer.lex(name, glueType);
     }
@@ -275,15 +280,17 @@ protected Field convertField(String name, String glueType)
     public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, GetTableRequest request) throws Exception
     {
         Map<String, String> qptArguments = request.getQueryPassthroughArguments();
-        queryPassthrough.verify(qptArguments);
-        String schemaName = qptArguments.get(NeptuneQueryPassthrough.DATABASE);
-        String tableName = qptArguments.get(NeptuneQueryPassthrough.COLLECTION);
-        TableName tableNameObj = new TableName(schemaName, tableName);
-        request = new GetTableRequest(request.getIdentity(), request.getQueryId(),
-                request.getCatalogName(), tableNameObj, request.getQueryPassthroughArguments());
+        if (qptArguments.containsKey(NeptuneGremlinQueryPassthrough.TRAVERSE)) {
+            gremlinQueryPassthrough.verify(qptArguments);
+        }
+        else {
+            sparqlQueryPassthrough.verify(qptArguments);
+        }
 
-        GetTableResponse getTableResponse = doGetTable(allocator, request);
-        List<Field> fields = getTableResponse.getSchema().getFields();
+        String schemaName;
+        String tableName;
+        String componentTypeValue;
+        TableName tableNameObj;
         Schema schema;
         Enums.GraphType graphType = Enums.GraphType.PROPERTYGRAPH;
 
@@ -293,19 +300,23 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
 
         switch (graphType){
             case PROPERTYGRAPH:
+                schemaName = qptArguments.get(NeptuneGremlinQueryPassthrough.DATABASE);
+                tableName = qptArguments.get(NeptuneGremlinQueryPassthrough.COLLECTION);
+                componentTypeValue = qptArguments.get(NeptuneGremlinQueryPassthrough.COMPONENT_TYPE);
+                tableNameObj = new TableName(schemaName, tableName);
                 Client client = neptuneConnection.getNeptuneClientConnection();
                 GraphTraversalSource graphTraversalSource = neptuneConnection.getTraversalSource(client);
-                String gremlinQuery = qptArguments.get(NeptuneQueryPassthrough.QUERY);
-                gremlinQuery = gremlinQuery.concat(".limit(1)");
+                String gremlinQuery = qptArguments.get(NeptuneGremlinQueryPassthrough.TRAVERSE);
+                gremlinQuery = gremlinQuery.concat("." + GREMLIN_QUERY_LIMIT);
                 logger.info("NeptuneMetadataHandler doGetQueryPassthroughSchema gremlinQuery with limit: " + gremlinQuery);
                 Object object = new PropertyGraphHandler(neptuneConnection).getResponseFromGremlinQuery(graphTraversalSource, gremlinQuery);
                 GraphTraversal graphTraversalForSchema = (GraphTraversal) object;
                 if (graphTraversalForSchema.hasNext()) {
                     Object responseObj = graphTraversalForSchema.next();
-                    if (responseObj instanceof Map && gremlinQuery.contains(Constants.GREMLIN_QUERY_SUPPORT_TYPE)) {
+                    if (responseObj instanceof Map) {
                         logger.info("NeptuneMetadataHandler doGetQueryPassthroughSchema gremlinQuery with valueMap");
                         Map graphTraversalObj = (Map) responseObj;
-                        schema = getSchemaFromResults(getTableResponse, graphTraversalObj, fields);
+                        schema = NeptuneSchemaUtils.getSchemaFromResults(graphTraversalObj, componentTypeValue, tableName);
                         return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
                     }
                     else {
@@ -319,16 +330,17 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
                 }
 
             case RDF:
-                String sparqlQuery = qptArguments.get(NeptuneQueryPassthrough.QUERY);
-                sparqlQuery = sparqlQuery.contains("limit") ? sparqlQuery : sparqlQuery.concat("\nlimit 1");
+                schemaName = qptArguments.get(NeptuneSparqlQueryPassthrough.DATABASE);
+                tableName = qptArguments.get(NeptuneSparqlQueryPassthrough.COLLECTION);
+                tableNameObj = new TableName(schemaName, tableName);
+                String sparqlQuery = qptArguments.get(NeptuneSparqlQueryPassthrough.QUERY);
+                sparqlQuery = sparqlQuery.contains("limit") ? sparqlQuery : sparqlQuery.concat("\n" + SPARQL_QUERY_LIMIT);
                 logger.info("NeptuneMetadataHandler doGetQueryPassthroughSchema sparql query with limit: " + sparqlQuery);
                 NeptuneSparqlConnection neptuneSparqlConnection = (NeptuneSparqlConnection) neptuneConnection;
                 neptuneSparqlConnection.runQuery(sparqlQuery);
-                String strim = getTableResponse.getSchema().getCustomMetadata().get(Constants.SCHEMA_STRIP_URI);
-                boolean trimURI = strim == null ? false : Boolean.parseBoolean(strim);
                 if (neptuneSparqlConnection.hasNext()) {
-                    Map<String, Object> resultsMap = neptuneSparqlConnection.next(trimURI);
-                    schema = getSchemaFromResults(getTableResponse, resultsMap, fields);
+                    Map<String, Object> resultsMap = neptuneSparqlConnection.next();
+                    schema = NeptuneSchemaUtils.getSchemaFromResults(resultsMap, RDF_COMPONENT_TYPE, tableName);
                     return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
                 }
                 else {
@@ -339,36 +351,4 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
                 throw new IllegalArgumentException("Unsupported graphType: " + graphType);
         }
     }
-
-    private Schema getSchemaFromResults(GetTableResponse getTableResponse, Map resultsMap, List<Field> fields)
-    {
-        Schema schema;
-        //In case of 'gremlin/sparql query' is fetching all columns then we can use same schema from glue
-        //otherwise we will build schema from gremlin/sparql query result column names.
-        if (resultsMap != null && resultsMap.size() == fields.size()) {
-            schema = getTableResponse.getSchema();
-        }
-        else {
-            SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
-            //Building schema from gremlin/sparql query results and list of fields from glue response.
-            //It's require only when we are selecting limited columns.
-            resultsMap.forEach((columnName, columnValue) -> buildSchema(columnName.toString(), fields, schemaBuilder));
-            Map<String, String> metaData = getTableResponse.getSchema().getCustomMetadata();
-            for (Map.Entry<String, String> map : metaData.entrySet()) {
-                schemaBuilder.addMetadata(map.getKey(), map.getValue());
-            }
-            schema = schemaBuilder.build();
-        }
-        return schema;
-    }
-
-    private void buildSchema(String columnName, List<Field> fields, SchemaBuilder schemaBuilder)
-    {
-        for (Field field : fields) {
-            if (field.getName().equals(columnName)) {
-                schemaBuilder.addField(field);
-                break;
-            }
-        }
-    }
 }

athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneSchemaUtils.java

@@ -0,0 +1,101 @@
+/*-
+ * #%L
+ * athena-neptune
+ * %%
+ * Copyright (C) 2019 - 2020 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.connectors.neptune;
+
+import com.amazonaws.athena.connector.lambda.data.SchemaBuilder;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.math.BigInteger;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Collection of helpful utilities that handle Neptune schema inference, type, and naming conversion.
+ */
+public class NeptuneSchemaUtils
+{
+    private static final Logger logger = LoggerFactory.getLogger(NeptuneSchemaUtils.class);
+
+    private NeptuneSchemaUtils() {}
+
+    public static Schema getSchemaFromResults(Map resultsMap, String componentTypeValue, String tableName)
+    {
+        Schema schema;
+        SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
+        //Building schema from gremlin/sparql query results.
+        resultsMap.forEach((columnName, columnValue) -> buildSchema(columnName.toString(), columnValue,  schemaBuilder));
+        schemaBuilder.addMetadata(Constants.SCHEMA_COMPONENT_TYPE, componentTypeValue);
+        schemaBuilder.addMetadata(Constants.SCHEMA_GLABEL, tableName);
+        schema = schemaBuilder.build();
+        return schema;
+    }
+
+    private static void buildSchema(String columnName, Object columnValue, SchemaBuilder schemaBuilder)
+    {
+        schemaBuilder.addField(getArrowFieldForNeptune(columnName, columnValue));
+    }
+
+    /**
+     * Infers the type of a field from Neptune data.
+     *
+     * @param key The key of the field we are attempting to infer.
+     * @param value A value from the key whose type we are attempting to infer.
+     * @return The Apache Arrow field definition of the inferred key/value.
+     */
+    private static Field getArrowFieldForNeptune(String key, Object value)
+    {
+        if (value instanceof String || value instanceof java.util.UUID) {
+            return new Field(key, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null);
+        }
+        else if (value instanceof Integer) {
+            return new Field(key, FieldType.nullable(Types.MinorType.INT.getType()), null);
+        }
+        else if (value instanceof BigInteger) {
+            return new Field(key, FieldType.nullable(Types.MinorType.BIGINT.getType()), null);
+        }
+        else if (value instanceof Long) {
+            return new Field(key, FieldType.nullable(Types.MinorType.BIGINT.getType()), null);
+        }
+        else if (value instanceof Boolean) {
+            return new Field(key, FieldType.nullable(Types.MinorType.BIT.getType()), null);
+        }
+        else if (value instanceof Float) {
+            return new Field(key, FieldType.nullable(Types.MinorType.FLOAT4.getType()), null);
+        }
+        else if (value instanceof Double) {
+            return new Field(key, FieldType.nullable(Types.MinorType.FLOAT8.getType()), null);
+        }
+        else if (value instanceof java.util.Date) {
+            return new Field(key, FieldType.nullable(Types.MinorType.DATEMILLI.getType()), null);
+        }
+        else if (value instanceof List) {
+            return getArrowFieldForNeptune(key, ((List<?>) value).get(0));
+        }
+
+        String className = (value == null || value.getClass() == null) ? "null" : value.getClass().getName();
+        logger.warn("Unknown type[{}] for field[{}], defaulting to varchar.", className, key);
+        return new Field(key, FieldType.nullable(Types.MinorType.VARCHAR.getType()), null);
+    }
+}

athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/propertygraph/PropertyGraphHandler.java

@@ -30,7 +30,7 @@
 import com.amazonaws.athena.connectors.neptune.propertygraph.rowwriters.CustomSchemaRowWriter;
 import com.amazonaws.athena.connectors.neptune.propertygraph.rowwriters.EdgeRowWriter;
 import com.amazonaws.athena.connectors.neptune.propertygraph.rowwriters.VertexRowWriter;
-import com.amazonaws.athena.connectors.neptune.qpt.NeptuneQueryPassthrough;
+import com.amazonaws.athena.connectors.neptune.qpt.NeptuneGremlinQueryPassthrough;
 import org.apache.arrow.util.VisibleForTesting;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.tinkerpop.gremlin.driver.Client;
@@ -73,7 +73,7 @@ public class PropertyGraphHandler
      */
 
     private final NeptuneConnection neptuneConnection;
-    private final NeptuneQueryPassthrough queryPassthrough = new NeptuneQueryPassthrough();
+    private final NeptuneGremlinQueryPassthrough queryPassthrough = new NeptuneGremlinQueryPassthrough();
 
     @VisibleForTesting
     public PropertyGraphHandler(NeptuneConnection neptuneConnection) 
@@ -117,8 +117,8 @@ public void executeQuery(
         if (recordsRequest.getConstraints().isQueryPassThrough()) {
             Map<String, String> qptArguments = recordsRequest.getConstraints().getQueryPassthroughArguments();
             queryPassthrough.verify(qptArguments);
-            labelName = qptArguments.get(NeptuneQueryPassthrough.COLLECTION);
-            gremlinQuery = qptArguments.get(NeptuneQueryPassthrough.QUERY);
+            labelName = qptArguments.get(NeptuneGremlinQueryPassthrough.COLLECTION);
+            gremlinQuery = qptArguments.get(NeptuneGremlinQueryPassthrough.TRAVERSE);
 
             Object object = getResponseFromGremlinQuery(graphTraversalSource, gremlinQuery);
             graphTraversal = (GraphTraversal) object;