Base array functionality working with IT tests.

Signed-off-by: forestmvey <[email protected]>
Bit-Quill · Jul 13, 2023 · 41d7bf5 · 41d7bf5
1 parent 579d44a
commit 41d7bf5
Show file tree

Hide file tree

Showing 8 changed files with 118 additions and 21 deletions.
diff --git a/common/src/main/java/org/opensearch/sql/common/utils/StringUtils.java b/common/src/main/java/org/opensearch/sql/common/utils/StringUtils.java
@@ -107,6 +107,6 @@ private static boolean isQuoted(String text, String mark) {
   }
 
   public static String removeParenthesis(String qualifier) {
-    return qualifier.replaceAll("\\[.+\\]", "");
+    return qualifier.replaceAll("\\[\\d+\\]", "");
   }
 }
diff --git a/core/src/main/java/org/opensearch/sql/expression/ArrayReferenceExpression.java b/core/src/main/java/org/opensearch/sql/expression/ArrayReferenceExpression.java
@@ -17,6 +17,8 @@
 import lombok.Getter;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.common.utils.StringUtils;
+import org.opensearch.sql.data.model.ExprCollectionValue;
+import org.opensearch.sql.data.model.ExprMissingValue;
 import org.opensearch.sql.data.model.ExprTupleValue;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -31,13 +33,13 @@ public class ArrayReferenceExpression extends ReferenceExpression {
   @Getter
   private final ExprType type;
   public ArrayReferenceExpression(String ref, ExprType type, List<Pair<String, OptionalInt>> partsAndIndexes) {
-    super(ref, type);
+    super(StringUtils.removeParenthesis(ref), type);
     this.partsAndIndexes = partsAndIndexes;
     this.type = type;
   }
 
   public ArrayReferenceExpression(ReferenceExpression ref) {
-    super(ref.toString(), ref.type());
+    super(StringUtils.removeParenthesis(ref.toString()), ref.type());
     this.partsAndIndexes = Arrays.stream(ref.toString().split("\\.")).map(e -> Pair.of(e, OptionalInt.empty())).collect(
         Collectors.toList());
     this.type = ref.type();
@@ -68,28 +70,25 @@ private ExprValue resolve(ExprValue value, List<Pair<String, OptionalInt>> paths
     ExprValue wholePathValue = value.keyValue(String.join(PATH_SEP, pathsWithoutParenthesis));
 
     if (!paths.get(0).getRight().isEmpty()) {
-      wholePathValue = getIndexValueOfCollection(wholePathValue);
-    } else if (paths.size() == 1 && !wholePathValue.type().equals(ExprCoreType.ARRAY)) {
-      wholePathValue = ExprValueUtils.missingValue();
+      if (value.keyValue(pathsWithoutParenthesis.get(0)) instanceof ExprCollectionValue) { // TODO check array size
+        wholePathValue = value
+            .keyValue(pathsWithoutParenthesis.get(0))
+            .collectionValue()
+            .get(paths.get(0).getRight().getAsInt());
+        if (paths.size() != 1) {
+          return resolve(wholePathValue, paths.subList(1, paths.size()));
+        }
+      } else {
+        return ExprValueUtils.missingValue();
+      }
+    } else if (wholePathValue.isMissing()) {
+      return resolve(value.keyValue(pathsWithoutParenthesis.get(0)), paths.subList(1, paths.size()));
     }
 
     if (!wholePathValue.isMissing() || paths.size() == 1) {
       return wholePathValue;
     } else {
-      return resolve(value.keyValue(pathsWithoutParenthesis.get(0)
-           ), paths.subList(1, paths.size()));
+      return resolve(value.keyValue(pathsWithoutParenthesis.get(0)), paths.subList(1, paths.size()));
     }
   }
-
-  private ExprValue getIndexValueOfCollection(ExprValue value) {
-    ExprValue collectionVal = value;
-//    for (OptionalInt currentIndex : List.of(index)) {
-      if (collectionVal.type().equals(ExprCoreType.ARRAY)) {
-//        collectionVal = collectionVal.collectionValue().get(currentIndex.getAsInt());
-      } else {
-        return ExprValueUtils.missingValue();
-      }
-//    }
-    return collectionVal;
-  }
 }
diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java
@@ -38,6 +38,7 @@
 import static com.google.common.base.Strings.isNullOrEmpty;
 import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient;
 import static org.opensearch.sql.legacy.TestUtils.getAccountIndexMapping;
+import static org.opensearch.sql.legacy.TestUtils.getArraysIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getBankIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getBankWithNullValuesIndexMapping;
 import static org.opensearch.sql.legacy.TestUtils.getDataTypeNonnumericIndexMapping;
@@ -681,7 +682,11 @@ public enum Index {
     NESTED_WITH_NULLS(TestsConstants.TEST_INDEX_NESTED_WITH_NULLS,
         "multi_nested",
         getNestedTypeIndexMapping(),
-        "src/test/resources/nested_with_nulls.json");
+        "src/test/resources/nested_with_nulls.json"),
+    ARRAYS(TestsConstants.TEST_INDEX_ARRAYS,
+        "arrays",
+        getArraysIndexMapping(),
+        "src/test/resources/arrays.json");
 
     private final String name;
     private final String type;

diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/TestUtils.java b/integ-test/src/test/java/org/opensearch/sql/legacy/TestUtils.java
@@ -223,6 +223,11 @@ public static String getDateIndexMapping() {
     return getMappingFile(mappingFile);
   }
 
+  public static String getArraysIndexMapping() {
+    String mappingFile = "arrays_mapping.json";
+    return getMappingFile(mappingFile);
+  }
+
   public static String getDateTimeIndexMapping() {
     String mappingFile = "date_time_index_mapping.json";
     return getMappingFile(mappingFile);

diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/TestsConstants.java b/integ-test/src/test/java/org/opensearch/sql/legacy/TestsConstants.java
@@ -48,6 +48,7 @@ public class TestsConstants {
   public final static String TEST_INDEX_ORDER = TEST_INDEX + "_order";
   public final static String TEST_INDEX_WEBLOG = TEST_INDEX + "_weblog";
   public final static String TEST_INDEX_DATE = TEST_INDEX + "_date";
+  public final static String TEST_INDEX_ARRAYS = TEST_INDEX + "_arrays";
   public final static String TEST_INDEX_DATE_TIME = TEST_INDEX + "_datetime";
   public final static String TEST_INDEX_DEEP_NESTED = TEST_INDEX + "_deep_nested";
   public final static String TEST_INDEX_STRINGS = TEST_INDEX + "_strings";

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/ArraysIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/ArraysIT.java
@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.sql;
+
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ARRAYS;
+import static org.opensearch.sql.util.MatcherUtils.rows;
+import static org.opensearch.sql.util.MatcherUtils.verifyDataRows;
+
+import java.io.IOException;
+import java.util.List;
+
+import com.google.common.collect.ImmutableMap;
+import org.json.JSONArray;
+import org.json.JSONObject;
+import org.junit.Test;
+import org.opensearch.sql.legacy.SQLIntegTestCase;
+
+public class ArraysIT extends SQLIntegTestCase {
+  @Override
+  public void init() throws IOException {
+    loadIndex(Index.ARRAYS);
+  }
+
+  @Test
+  public void object_array_index_1_test() {
+    String query = "SELECT objectArray[0] FROM " + TEST_INDEX_ARRAYS;
+    JSONObject result = executeJdbcRequest(query);
+
+    verifyDataRows(result,
+        rows(new JSONObject(ImmutableMap.of("innerObject", List.of(1, 2)))));
+  }
+
+  @Test
+  public void object_array_index_1_inner_object_test() {
+    String query = "SELECT objectArray[0].innerObject FROM " + TEST_INDEX_ARRAYS;
+    JSONObject result = executeJdbcRequest(query);
+
+    verifyDataRows(result,
+        rows(new JSONArray(List.of(1, 2))));
+  }
+
+  @Test
+  public void object_array_index_1_inner_object_index_1_test() {
+    String query = "SELECT objectArray[0].innerObject[0] FROM " + TEST_INDEX_ARRAYS;
+    JSONObject result = executeJdbcRequest(query);
+
+    verifyDataRows(result,
+        rows(1));
+  }
+
+  @Test
+  public void multi_object_array_index_1_test() {
+    String query = "SELECT multiObjectArray[0] FROM " + TEST_INDEX_ARRAYS;
+    JSONObject result = executeJdbcRequest(query);
+
+    verifyDataRows(result,
+        rows(new JSONObject(ImmutableMap.of("id", 1, "name", "blah"))));
+  }
+}
diff --git a/integ-test/src/test/resources/arrays.json b/integ-test/src/test/resources/arrays.json
@@ -0,0 +1,2 @@
+{"index":{"_id":"1"}}
+{"objectArray": [{"innerObject": [1, 2]}, {"innerObject": 3}], "multiObjectArray":[{"id": 1, "name": "blah"}, {"id": 2,"name": "hoo"}]}
diff --git a/integ-test/src/test/resources/indexDefinitions/arrays_mapping.json b/integ-test/src/test/resources/indexDefinitions/arrays_mapping.json
@@ -0,0 +1,23 @@
+{
+  "mappings": {
+    "properties": {
+      "objectArray": {
+        "properties": {
+          "innerObject": {
+            "type": "keyword"
+          }
+        }
+      },
+      "multiObjectArray": {
+        "properties": {
+          "id": {
+            "type": "long"
+          },
+          "name": {
+            "type": "keyword"
+          }
+        }
+      }
+    }
+  }
+}