Athena to Amazon Neptune connector code

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

@@ -22,17 +22,18 @@
 
 import com.amazonaws.athena.connector.lambda.domain.predicate.Marker.Bound;
 
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.tinkerpop.gremlin.process.traversal.P;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversal;
-import org.apache.tinkerpop.gremlin.process.traversal.lambda.PredicateTraversal;
 import org.apache.tinkerpop.gremlin.structure.Vertex;
 
 /**
  * This class is a Utility class to general gremlin query equivalents of
  * Contraints being passed via AWS Lambda Handler
  */
 public class GremlinQueryPreProcessor {
-    
+
     public enum Operator {
         LESSTHAN, GREATERTHAN, EQUALTO, NOTEQUALTO
     }
@@ -49,17 +50,13 @@ public enum Operator {
      * 
      * @return A Gremlin Query Part equivalent to Contraint.
      */
-    // public static String generateGremlinQueryPart(String key, String value,
-    // String type, Bound bound,
-    // Operator operator) {
-
     public static GraphTraversal<Vertex, Vertex> generateGremlinQueryPart(GraphTraversal<Vertex, Vertex> traversal,
-            String key, String value, String type, Bound bound, Operator operator) {
-
-        switch (type) {
+            String key, String value, ArrowType type, Bound bound, Operator operator) {
 
-            case "Int(32, true)":
+        Types.MinorType minorType = Types.getMinorTypeForArrowType(type);
 
+        switch (minorType) {
+            case INT:
                 if (operator.equals(Operator.GREATERTHAN)) {
 
                     traversal = bound.equals(Bound.EXACTLY) ? traversal.has(key, P.gte(Integer.parseInt(value)))
@@ -84,7 +81,59 @@ public static GraphTraversal<Vertex, Vertex> generateGremlinQueryPart(GraphTrave
 
                 break;
 
-            case "Utf8":
+            case FLOAT4:
+
+                if (operator.equals(Operator.GREATERTHAN)) {
+
+                    traversal = bound.equals(Bound.EXACTLY) ? traversal.has(key, P.gte(Float.parseFloat(value)))
+                            : traversal.has(key, P.gt(Float.parseFloat(value)));
+                }
+
+                if (operator.equals(Operator.LESSTHAN)) {
+
+                    traversal = bound.equals(Bound.EXACTLY) ? traversal.has(key, P.lte(Float.parseFloat(value)))
+                            : traversal.has(key, P.lt(Float.parseFloat(value)));
+                }
+
+                if (operator.equals(Operator.EQUALTO)) {
+
+                    traversal = traversal.has(key, P.eq(Float.parseFloat(value)));
+                }
+
+                if (operator.equals(Operator.NOTEQUALTO)) {
+
+                    traversal = traversal.has(key, P.neq(Float.parseFloat(value)));
+                }
+
+                break;
+
+            case FLOAT8:
+
+                if (operator.equals(Operator.GREATERTHAN)) {
+
+                    traversal = bound.equals(Bound.EXACTLY) ? traversal.has(key, P.gte(Double.parseDouble(value)))
+                            : traversal.has(key, P.gt(Double.parseDouble(value)));
+                }
+
+                if (operator.equals(Operator.LESSTHAN)) {
+
+                    traversal = bound.equals(Bound.EXACTLY) ? traversal.has(key, P.lte(Double.parseDouble(value)))
+                            : traversal.has(key, P.lt(Double.parseDouble(value)));
+                }
+
+                if (operator.equals(Operator.EQUALTO)) {
+
+                    traversal = traversal.has(key, P.eq(Double.parseDouble(value)));
+                }
+
+                if (operator.equals(Operator.NOTEQUALTO)) {
+
+                    traversal = traversal.has(key, P.neq(Integer.parseInt(value)));
+                }
+
+                break;
+
+            case VARCHAR:
 
                 if (operator.equals(Operator.EQUALTO)) {
 

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

@@ -173,9 +173,12 @@ protected void readWithConstraint(final BlockSpiller spiller, final ReadRecordsR
 
     /**
      * Used to generate Gremlin Query part for Constraint Map
-     *
+     * 
+     * @param traversal Gremlin Traversal, traversal is updated based on constraints map
      * @param hasMap Constraint Hash Map
-     **/
+     * 
+     * @return A Gremlin Query Part equivalent to Contraint.
+     */
     public GraphTraversal<Vertex, Vertex> getQueryPartForContraintsMap(GraphTraversal<Vertex, Vertex> traversal,
             final Map hashMap) {
 
@@ -190,7 +193,7 @@ public GraphTraversal<Vertex, Vertex> getQueryPartForContraintsMap(GraphTraversa
                     if (range.getLow().getValue().toString().equals(range.getHigh().getValue().toString())) {
 
                         traversal = GremlinQueryPreProcessor.generateGremlinQueryPart(traversal, key,
-                                range.getLow().getValue().toString(), range.getType().toString(),
+                                range.getLow().getValue().toString(), range.getType(),
                                 range.getLow().getBound(), GremlinQueryPreProcessor.Operator.EQUALTO);
                         break;
                     }
@@ -202,14 +205,14 @@ public GraphTraversal<Vertex, Vertex> getQueryPartForContraintsMap(GraphTraversa
                             + range.getType().toString().equalsIgnoreCase(Types.MinorType.INT.getType().toString()));
 
                     traversal = GremlinQueryPreProcessor.generateGremlinQueryPart(traversal, key,
-                            range.getLow().getValue().toString(), range.getType().toString(), range.getLow().getBound(),
+                            range.getLow().getValue().toString(), range.getType(), range.getLow().getBound(),
                             GremlinQueryPreProcessor.Operator.GREATERTHAN);
                 }
 
                 if (!range.getHigh().isNullValue()) {
 
                     traversal = GremlinQueryPreProcessor.generateGremlinQueryPart(traversal, key,
-                            range.getHigh().getValue().toString(), range.getType().toString(),
+                            range.getHigh().getValue().toString(), range.getType(),
                             range.getLow().getBound(), GremlinQueryPreProcessor.Operator.LESSTHAN);
                 }
             }

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

@@ -28,19 +28,18 @@
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import com.amazonaws.athena.connector.lambda.data.writers.holders.NullableVarCharHolder;
 import org.apache.arrow.vector.holders.NullableIntHolder;
 import org.apache.arrow.vector.holders.NullableFloat4Holder;
 import org.apache.arrow.vector.holders.NullableFloat8Holder;
 
+/**
+ * This class is a Utility class to create Extractors for each fields as per Schema
+ */
 public class TypeRowWriter {
 
     public static RowWriterBuilder writeRowTemplate(RowWriterBuilder rowWriterBuilder, Field field) {
 
-        Logger logger = LoggerFactory.getLogger(TypeRowWriter.class);
-
         ArrowType arrowType = field.getType();
         Types.MinorType minorType = Types.getMinorTypeForArrowType(arrowType);
 
@@ -50,11 +49,10 @@ public static RowWriterBuilder writeRowTemplate(RowWriterBuilder rowWriterBuilde
                 rowWriterBuilder.withExtractor(field.getName(),
                         (Float8Extractor) (Object context, NullableFloat8Holder value) -> {
                             value.isSet = 1;
-
                             Map<Object, Object> obj = (Map<Object, Object>) context;
                             ArrayList<Object> objValues = (ArrayList)obj.get(field.getName());
 
-                            value.value = Float.parseFloat(objValues.get(0).toString());
+                            value.value = Double.parseDouble(objValues.get(0).toString());
                         });
 
                 break;
@@ -76,7 +74,6 @@ public static RowWriterBuilder writeRowTemplate(RowWriterBuilder rowWriterBuilde
                 rowWriterBuilder.withExtractor(field.getName(),
                         (Float4Extractor) (Object context, NullableFloat4Holder value) -> {
                             value.isSet = 1;
-
                             Map<Object, Object> obj = (Map<Object, Object>) context;
                             ArrayList<Object> objValues = (ArrayList)obj.get(field.getName());
 
@@ -88,7 +85,6 @@ public static RowWriterBuilder writeRowTemplate(RowWriterBuilder rowWriterBuilde
                 rowWriterBuilder.withExtractor(field.getName(),
                         (IntExtractor) (Object context, NullableIntHolder value) -> {
                             value.isSet = 1;
-
                             Map<Object, Object> obj = (Map<Object, Object>) context;
                             ArrayList<Object> objValues = (ArrayList)obj.get(field.getName());
                             value.value = Integer.parseInt(objValues.get(0).toString());

athena-neptune/src/test/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandlerTest.java

@@ -19,33 +19,14 @@
  */
 package com.amazonaws.athena.connectors.neptune;
 
-import com.amazonaws.athena.connector.lambda.QueryStatusChecker;
-import com.amazonaws.athena.connector.lambda.data.Block;
-import com.amazonaws.athena.connector.lambda.data.BlockAllocator;
 import com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl;
-import com.amazonaws.athena.connector.lambda.data.BlockUtils;
-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.predicate.Constraints;
-import com.amazonaws.athena.connector.lambda.domain.predicate.Range;
-import com.amazonaws.athena.connector.lambda.domain.predicate.SortedRangeSet;
-import com.amazonaws.athena.connector.lambda.domain.predicate.ValueSet;
-import com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler;
-import com.amazonaws.athena.connector.lambda.metadata.GetSplitsRequest;
-import com.amazonaws.athena.connector.lambda.metadata.GetSplitsResponse;
-import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutRequest;
-import com.amazonaws.athena.connector.lambda.metadata.GetTableLayoutResponse;
 import com.amazonaws.athena.connector.lambda.metadata.GetTableRequest;
 import com.amazonaws.athena.connector.lambda.metadata.GetTableResponse;
 import com.amazonaws.athena.connector.lambda.metadata.ListSchemasRequest;
 import com.amazonaws.athena.connector.lambda.metadata.ListSchemasResponse;
 import com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest;
 import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
-import com.amazonaws.athena.connector.lambda.metadata.MetadataRequestType;
-import com.amazonaws.athena.connector.lambda.metadata.MetadataResponse;
-import com.amazonaws.athena.connector.lambda.security.FederatedIdentity;
 import com.amazonaws.athena.connector.lambda.security.LocalKeyFactory;
 import com.amazonaws.services.athena.AmazonAthena;
 import com.amazonaws.services.glue.AWSGlue;
@@ -56,10 +37,7 @@
 import com.amazonaws.services.glue.model.StorageDescriptor;
 import com.amazonaws.services.glue.model.Table;
 import com.amazonaws.services.secretsmanager.AWSSecretsManager;
-import com.google.common.collect.ImmutableList;
-import org.apache.arrow.vector.types.Types;
-import org.apache.arrow.vector.types.pojo.Schema;
-import org.apache.commons.collections.CursorableLinkedList.Cursor;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -70,20 +48,15 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 import static org.junit.Assert.*;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
-import org.mockito.invocation.InvocationOnMock;
 import org.mockito.runners.MockitoJUnitRunner;
 
 @RunWith(MockitoJUnitRunner.class)