DocDB select query with _id issue fix (#2105)

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

athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/QueryUtils.java

@@ -45,10 +45,12 @@
 import org.apache.arrow.vector.util.Text;
 import org.bson.Document;
 import org.bson.json.JsonParseException;
+import org.bson.types.ObjectId;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Verify.verify;
@@ -76,6 +78,7 @@ public final class QueryUtils
     private static final String LTE_OP = "$lte";
     private static final String IN_OP = "$in";
     private static final String NOTIN_OP = "$nin";
+    private static final String COLUMN_NAME_ID = "_id";
 
     private QueryUtils()
     {
@@ -204,10 +207,25 @@ public static Document makePredicate(Field field, ValueSet constraint)
 
         // Add back all of the possible single values either as an equality or an IN predicate
         if (singleValues.size() == 1) {
-            disjuncts.add(documentOf(EQ_OP, singleValues.get(0)));
+            Object value = singleValues.get(0);
+            if (name.equals(COLUMN_NAME_ID)) {
+                ObjectId objectId = new ObjectId(value.toString());
+                disjuncts.add(documentOf(EQ_OP, objectId));
+            }
+            else {
+                disjuncts.add(documentOf(EQ_OP, value));
+            }
         }
         else if (singleValues.size() > 1) {
-            disjuncts.add(documentOf(IN_OP, singleValues));
+            if (name.equals(COLUMN_NAME_ID)) {
+                List<ObjectId> objectIdList = singleValues.stream()
+                        .map(obj -> new ObjectId(obj.toString()))
+                        .collect(Collectors.toList());
+                disjuncts.add(documentOf(IN_OP, objectIdList));
+            }
+            else {
+                disjuncts.add(documentOf(IN_OP, singleValues));
+            }
         }
 
         return orPredicate(disjuncts.stream()

athena-docdb/src/test/java/com/amazonaws/athena/connectors/docdb/QueryUtilsTest.java

@@ -0,0 +1,106 @@
+/*-
+ * #%L
+ * athena-docdb
+ * %%
+ * Copyright (C) 2019 - 2024 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.docdb;
+
+import com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl;
+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.google.common.collect.ImmutableList;
+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.apache.arrow.vector.types.pojo.FieldType;
+import org.bson.Document;
+import org.bson.types.ObjectId;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class QueryUtilsTest
+{
+    private final BlockAllocatorImpl allocator = new BlockAllocatorImpl();
+
+    @Test
+    public void testMakePredicateWithSortedRangeSet()
+    {
+        Field field = new Field("year", FieldType.nullable(new ArrowType.Int(32, true)), null);
+
+        ValueSet rangeSet = SortedRangeSet.copyOf(
+                Types.MinorType.INT.getType(),
+                ImmutableList.of(
+                        Range.lessThan(allocator, Types.MinorType.INT.getType(), 1950),
+                        Range.equal(allocator, Types.MinorType.INT.getType(), 1952),
+                        Range.range(allocator, Types.MinorType.INT.getType(), 1955, false, 1972, true),
+                        Range.greaterThanOrEqual(allocator, Types.MinorType.INT.getType(), 2010)),
+                false
+        );
+        Document result = QueryUtils.makePredicate(field, rangeSet);
+        assertNotNull(result);
+        Document expected = new Document("$or", ImmutableList.of(
+                new Document("year", new Document("$lt", 1950)),
+                new Document("year", new Document("$gt", 1955).append("$lte", 1972)),
+                new Document("year", new Document("$gte", 2010)),
+                new Document("year", new Document("$eq", 1952))
+        ));
+        assertEquals(expected, result);
+    }
+
+    @Test
+    public void testMakePredicateWithId()
+    {
+        Field field = new Field("_id", FieldType.nullable(new ArrowType.Utf8()), null);
+
+        ValueSet rangeSet = SortedRangeSet.copyOf(
+                Types.MinorType.VARCHAR.getType(),
+                ImmutableList.of(
+                        Range.equal(allocator, Types.MinorType.VARCHAR.getType(), "4ecbe7f9e8c1c9092c000027")),
+                false
+        );
+        Document result = QueryUtils.makePredicate(field, rangeSet);
+        assertNotNull(result);
+        Document expected =
+                new Document("_id", new Document("$eq", new ObjectId("4ecbe7f9e8c1c9092c000027")));
+        assertEquals(expected, result);
+    }
+
+    @Test
+    public void testParseFilter()
+    {
+        String jsonFilter = "{ \"field\": { \"$eq\": \"value\" } }";
+
+        Document result = QueryUtils.parseFilter(jsonFilter);
+        assertNotNull(result);
+        assertEquals("value", ((Document) result.get("field")).get("$eq"));
+    }
+
+    @Test
+    public void testParseFilterInvalidJson()
+    {
+        String invalidJsonFilter = "{ field: { $eq: value } }";
+
+        assertThrows(IllegalArgumentException.class, () -> {
+            QueryUtils.parseFilter(invalidJsonFilter);
+        });
+    }
+}
+