Add _routing to SQL includes list (#277) (opensearch-project#1771)

* Add _routing to SQL includes list Signed-off-by: Andrew Carbonetto <andrewc@bitquilltech.com> * Update IT index Signed-off-by: Andrew Carbonetto <andrewc@bitquilltech.com> * Update doctest Signed-off-by: acarbonetto <andrewc@bitquilltech.com> * Add WHERE clause IT tests Signed-off-by: acarbonetto <andrewc@bitquilltech.com> * Fix IT test Signed-off-by: acarbonetto <andrewc@bitquilltech.com> * Update documentation for _routing Signed-off-by: acarbonetto <andrewc@bitquilltech.com> --------- Signed-off-by: Andrew Carbonetto <andrewc@bitquilltech.com> Signed-off-by: acarbonetto <andrewc@bitquilltech.com> Signed-off-by: Mitchell Gale <Mitchell.Gale@improving.com>
Bit-Quill · Jul 11, 2023 · ee37fe0 · ee37fe0
1 parent 371bd6c
commit ee37fe0
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 8 deletions.
diff --git a/docs/user/dql/basics.rst b/docs/user/dql/basics.rst
@@ -155,14 +155,17 @@ Result set:
 |  Nanette|   Bates|
 +---------+--------+
 
-One can also provide meta-field name(s) to retrieve reserved-fields (beginning with underscore) from OpenSearch documents. Meta-fields are not output
-from wildcard calls (`SELECT *`) and must be explicitly included to be returned.
+One can also provide meta-field name(s) to retrieve reserved-fields (beginning with underscore) from OpenSearch documents.  They may also be used
+in the query `WHERE` or `ORDER BY` clauses. Meta-fields are not output from wildcard calls (`SELECT *`) and must be explicitly included to be returned.
+
+Note: `_routing` is used differently in the `SELECT` and `WHERE` clauses.  In `WHERE`, it contains the routing hash id. In `SELECT`,
+it returns the shard used for the query (unless shards aren't active, in which case it returns the routing hash id).
 
 SQL query::
 
 	POST /_plugins/_sql
 	{
-	  "query" : "SELECT firstname, lastname, _id, _index, _sort FROM accounts"
+	  "query" : "SELECT firstname, lastname, _id, _index, _sort, _routing FROM accounts WHERE _index = 'accounts'"
 	}
 
 Explain::
@@ -175,6 +178,7 @@ Explain::
 	      "firstname",
 	      "_id",
 	      "_index",
+	      "_routing",
 	      "_sort",
 	      "lastname"
 	    ],

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/IdentifierIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/IdentifierIT.java
@@ -14,6 +14,8 @@
 import static org.opensearch.sql.util.TestUtils.performRequest;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import org.json.JSONArray;
 import org.json.JSONObject;
 import org.junit.jupiter.api.Test;
 import org.opensearch.client.Request;
@@ -99,6 +101,80 @@ public void testMetafieldIdentifierTest() throws IOException {
     verifyDataRows(result, rows(30, id, index, 1.0, 1.0, -2));
   }
 
+  @Test
+  public void testMetafieldIdentifierRoutingSelectTest() throws IOException {
+    // create an index, but the contents doesn't really matter
+    String index = "test.routing_select";
+    String mapping = "{\"_routing\": {\"required\": true }}";
+    new Index(index, mapping)
+        .addDocWithShardId("{\"age\": 31}", "test0", "test0")
+        .addDocWithShardId("{\"age\": 31}", "test1", "test1")
+        .addDocWithShardId("{\"age\": 32}", "test2", "test2")
+        .addDocWithShardId("{\"age\": 33}", "test3", "test3")
+        .addDocWithShardId("{\"age\": 34}", "test4", "test4")
+        .addDocWithShardId("{\"age\": 35}", "test5", "test5");
+
+    // Execute using field metadata values filtering on the routing shard hash id
+    final JSONObject result = new JSONObject(executeQuery(
+        "SELECT age, _id, _index, _routing "
+            + "FROM " + index,
+        "jdbc"));
+
+    // Verify that the metadata values are returned when requested
+    verifySchema(result,
+        schema("age", null, "long"),
+        schema("_id", null, "keyword"),
+        schema("_index", null, "keyword"),
+        schema("_routing", null, "keyword"));
+    assertTrue(result.getJSONArray("schema").length() == 4);
+
+    var datarows = result.getJSONArray("datarows");
+    assertEquals(6, datarows.length());
+
+    // note that _routing in the SELECT clause returns the shard
+    for (int i = 0; i < 6; i++) {
+      assertEquals("test" + i, datarows.getJSONArray(i).getString(1));
+      assertEquals(index, datarows.getJSONArray(i).getString(2));
+      assertTrue(datarows.getJSONArray(i).getString(3).contains("[" + index + "]"));
+    }
+  }
+
+  @Test
+  public void testMetafieldIdentifierRoutingFilterTest() throws IOException {
+    // create an index, but the contents doesn't really matter
+    String index = "test.routing_filter";
+    String mapping = "{\"_routing\": {\"required\": true }}";
+    new Index(index, mapping)
+        .addDocWithShardId("{\"age\": 31}", "test1", "test1")
+        .addDocWithShardId("{\"age\": 32}", "test2", "test2")
+        .addDocWithShardId("{\"age\": 33}", "test3", "test3")
+        .addDocWithShardId("{\"age\": 34}", "test4", "test4")
+        .addDocWithShardId("{\"age\": 35}", "test5", "test5")
+        .addDocWithShardId("{\"age\": 36}", "test6", "test6");
+
+    // Execute using field metadata values filtering on the routing shard hash id
+    final JSONObject result = new JSONObject(executeQuery(
+        "SELECT _id, _index, _routing "
+            + "FROM " + index + " "
+            + "WHERE _routing = \\\"test4\\\"",
+        "jdbc"));
+
+    // Verify that the metadata values are returned when requested
+    verifySchema(result,
+        schema("_id", null, "keyword"),
+        schema("_index", null, "keyword"),
+        schema("_routing", null, "keyword"));
+    assertTrue(result.getJSONArray("schema").length() == 3);
+
+    var datarows = result.getJSONArray("datarows");
+    assertEquals(1, datarows.length());
+
+    assertEquals("test4", datarows.getJSONArray(0).getString(0));
+    // note that _routing in the SELECT clause returns the shard, not the routing hash id
+    assertTrue(datarows.getJSONArray(0).getString(2).contains("[" + index + "]"));
+
+  }
+
   @Test
   public void testMetafieldIdentifierWithAliasTest() throws IOException {
     // create an index, but the contents doesn't matter
@@ -152,16 +228,32 @@ private static class Index {
       }
     }
 
+    Index(String indexName, String mapping) throws IOException {
+      this.indexName = indexName;
+
+      Request createIndex = new Request("PUT", "/" + indexName);
+      createIndex.setJsonEntity(mapping);
+      executeRequest(new Request("PUT", "/" + indexName));
+    }
+
     void addDoc(String doc) {
       Request indexDoc = new Request("POST", String.format("/%s/_doc?refresh=true", indexName));
       indexDoc.setJsonEntity(doc);
       performRequest(client(), indexDoc);
     }
 
-    void addDoc(String doc, String id) {
+    public Index addDoc(String doc, String id) {
       Request indexDoc = new Request("POST", String.format("/%s/_doc/%s?refresh=true", indexName, id));
       indexDoc.setJsonEntity(doc);
       performRequest(client(), indexDoc);
+      return this;
+    }
+
+    public Index addDocWithShardId(String doc, String id, String routing) {
+      Request indexDoc = new Request("POST", String.format("/%s/_doc/%s?refresh=true&routing=%s", indexName, id, routing));
+      indexDoc.setJsonEntity(doc);
+      performRequest(client(), indexDoc);
+      return this;
     }
   }
 

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java
@@ -10,6 +10,7 @@
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_INDEX;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_MAXSCORE;
+import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ROUTING;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_SCORE;
 import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_SORT;
 
@@ -185,8 +186,10 @@ private void addMetaDataFieldsToBuilder(
         if (maxScore != null) {
           builder.put(METADATA_FIELD_MAXSCORE, maxScore);
         }
-      } else { // if (metaDataField.equals(METADATA_FIELD_SORT)) {
+      } else if (metaDataField.equals(METADATA_FIELD_SORT)) {
         builder.put(METADATA_FIELD_SORT, new ExprLongValue(hit.getSeqNo()));
+      } else { // if (metaDataField.equals(METADATA_FIELD_ROUTING)){
+        builder.put(METADATA_FIELD_ROUTING, new ExprStringValue(hit.getShard().toString()));
       }
     });
   }

diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java
@@ -45,12 +45,15 @@ public class OpenSearchIndex implements Table {
   public static final String METADATA_FIELD_MAXSCORE = "_maxscore";
   public static final String METADATA_FIELD_SORT = "_sort";
 
+  public static final String METADATA_FIELD_ROUTING = "_routing";
+
   public static final java.util.Map<String, ExprType> METADATAFIELD_TYPE_MAP = Map.of(
       METADATA_FIELD_ID, ExprCoreType.STRING,
       METADATA_FIELD_INDEX, ExprCoreType.STRING,
       METADATA_FIELD_SCORE, ExprCoreType.FLOAT,
       METADATA_FIELD_MAXSCORE, ExprCoreType.FLOAT,
-      METADATA_FIELD_SORT, ExprCoreType.LONG
+      METADATA_FIELD_SORT, ExprCoreType.LONG,
+      METADATA_FIELD_ROUTING, ExprCoreType.STRING
   );
 
   /** OpenSearch client connection. */

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/response/OpenSearchResponseTest.java
@@ -30,8 +30,10 @@
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.common.bytes.BytesArray;
 import org.opensearch.common.text.Text;
+import org.opensearch.index.shard.ShardId;
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
+import org.opensearch.search.SearchShardTarget;
 import org.opensearch.search.aggregations.Aggregations;
 import org.opensearch.search.fetch.subphase.highlight.HighlightField;
 import org.opensearch.sql.data.model.ExprFloatValue;
@@ -146,9 +148,13 @@ void iterator_metafields() {
                 new TotalHits(1L, TotalHits.Relation.EQUAL_TO),
                 3.75F));
 
+    ShardId shardId = new ShardId("index", "indexUUID", 42);
+    SearchShardTarget shardTarget = new SearchShardTarget("node", shardId, null, null);
+
     when(searchHit1.getSourceAsString()).thenReturn("{\"id1\", 1}");
     when(searchHit1.getId()).thenReturn("testId");
     when(searchHit1.getIndex()).thenReturn("testIndex");
+    when(searchHit1.getShard()).thenReturn(shardTarget);
     when(searchHit1.getScore()).thenReturn(3.75F);
     when(searchHit1.getSeqNo()).thenReturn(123456L);
 
@@ -158,11 +164,12 @@ void iterator_metafields() {
         "id1", new ExprIntegerValue(1),
         "_index", new ExprStringValue("testIndex"),
         "_id", new ExprStringValue("testId"),
+        "_routing", new ExprStringValue(shardTarget.toString()),
         "_sort", new ExprLongValue(123456L),
         "_score", new ExprFloatValue(3.75F),
         "_maxscore", new ExprFloatValue(3.75F)
     ));
-    List includes = List.of("id1", "_index", "_id", "_sort", "_score", "_maxscore");
+    List includes = List.of("id1", "_index", "_id", "_routing", "_sort", "_score", "_maxscore");
     int i = 0;
     for (ExprValue hit : new OpenSearchResponse(searchResponse, factory, includes)) {
       if (i == 0) {

diff --git a/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java b/opensearch/src/test/java/org/opensearch/sql/opensearch/storage/OpenSearchIndexTest.java
@@ -187,9 +187,10 @@ void getReservedFieldTypes() {
     assertThat(
         fieldTypes,
         allOf(
-            aMapWithSize(5),
+            aMapWithSize(6),
             hasEntry("_id", ExprCoreType.STRING),
             hasEntry("_index", ExprCoreType.STRING),
+            hasEntry("_routing", ExprCoreType.STRING),
             hasEntry("_sort", ExprCoreType.LONG),
             hasEntry("_score", ExprCoreType.FLOAT),
             hasEntry("_maxscore", ExprCoreType.FLOAT)