Got the list and maps working.

Yelp · Mar 18, 2024 · c8ff362 · c8ff362
1 parent fca27c0
commit c8ff362
Show file tree

Hide file tree

Showing 5 changed files with 123 additions and 13 deletions.
diff --git a/clientlib/src/main/proto/yelp/nrtsearch/search.proto b/clientlib/src/main/proto/yelp/nrtsearch/search.proto
@@ -608,12 +608,30 @@ message SearchResponse {
                 google.protobuf.Struct structValue = 8; // Value for structured data
                 // Value for VECTOR FieldType
                 Vector vectorValue = 9;
+                RepeatedFieldValues repeatedFieldValues = 10;
             }
 
             message Vector {
                 repeated float value = 1;
             }
 
+            message RepeatedFieldValues {
+                // Repeated value for ATOM and TEXT  FieldType
+                repeated string textValues = 1; 
+                // Repeated value for BOOLEAN FieldType
+                repeated bool booleanValues = 2;
+                // Repeated value for INT FieldType2; 
+                repeated int32 intValues = 3; 
+                // Repeated value for LONG and DATE_TIME (as milliseconds since epoch in UTC zone) FieldType
+                repeated int64 longValues = 4; 
+                // Repeated value for FLOAT FieldType
+                repeated float floatValues = 5; 
+                // Repeated value for DOUBLE FieldType
+                repeated double doubleValues = 6; 
+                // Repeated value for structured data
+                repeated google.protobuf.Struct structValues = 7; 
+            }
+
         }
 
         message CompositeFieldValue {

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/SearchHandler.java
@@ -15,8 +15,11 @@
  */
 package com.yelp.nrtsearch.server.luceneserver;
 
+import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.protobuf.Struct;
+import com.google.protobuf.Struct.Builder;
+import com.google.protobuf.Value;
 import com.google.protobuf.util.JsonFormat;
 import com.yelp.nrtsearch.server.grpc.DeadlineUtils;
 import com.yelp.nrtsearch.server.grpc.FacetResult;
@@ -989,13 +992,108 @@ private static void fetchRuntimeFromSegmentFactory(
           } else if (obj instanceof Integer) {
             compositeFieldValue.addFieldValue(
                 SearchResponse.Hit.FieldValue.newBuilder().setIntValue((Integer) obj));
+          } else if (obj instanceof Map) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder()
+                    .setStructValue(createStruct((Map<String, Object>) obj)));
+          } else if (obj instanceof Boolean) {
+            compositeFieldValue.addFieldValue(
+                SearchResponse.Hit.FieldValue.newBuilder().setBooleanValue((Boolean) obj));
+          } else if (obj instanceof ArrayList) {
+            SearchResponse.Hit.FieldValue.Builder fieldValueBuilder =
+                SearchResponse.Hit.FieldValue.newBuilder();
+            SearchResponse.Hit.FieldValue.RepeatedFieldValues.Builder repeatedFieldValuesBuilder =
+                SearchResponse.Hit.FieldValue.RepeatedFieldValues.newBuilder();
+            generateRepeatedFields(repeatedFieldValuesBuilder, obj);
+            fieldValueBuilder.setRepeatedFieldValues(repeatedFieldValuesBuilder);
+            compositeFieldValue.addFieldValue(fieldValueBuilder);
           }
           // TODO: Add support for list and map.
           hit.putFields(name, compositeFieldValue.build());
         }
       }
     }
 
+    /**
+     * Generates list fields for all types except for lists.
+     *
+     * <p>TODO: Support list of lists.
+     *
+     * @param repeatedFieldValuesBuilder The builder for repeated field value in the search
+     *     response.
+     * @param objs Runtime field objects that are supposed to be list.
+     */
+    private static void generateRepeatedFields(
+        SearchResponse.Hit.FieldValue.RepeatedFieldValues.Builder repeatedFieldValuesBuilder,
+        Object objs) {
+      List<Struct> structObjs = new ArrayList<>();
+      for (Object obj : (Iterable<? extends Object>) objs) {
+        if (obj instanceof String) {
+          repeatedFieldValuesBuilder.addAllTextValues((Iterable<String>) objs);
+          return;
+        } else if (obj instanceof Integer) {
+          repeatedFieldValuesBuilder.addAllIntValues((Iterable<Integer>) objs);
+          return;
+        } else if (obj instanceof Double) {
+          repeatedFieldValuesBuilder.addAllDoubleValues((Iterable<Double>) objs);
+          return;
+        } else if (obj instanceof Float) {
+          repeatedFieldValuesBuilder.addAllFloatValues((Iterable<Float>) objs);
+          return;
+        } else if (obj instanceof Long) {
+          repeatedFieldValuesBuilder.addAllLongValues((Iterable<Long>) objs);
+          return;
+        } else if (obj instanceof Boolean) {
+          repeatedFieldValuesBuilder.addAllBooleanValues((Iterable<Boolean>) objs);
+          return;
+        } else if (obj instanceof Map) {
+          structObjs.add(createStruct((Map<String, Object>) obj));
+        } else {
+          return;
+        }
+      }
+
+      // Construct repeated structs.
+      if (Iterables.size(structObjs) > 0) {
+        repeatedFieldValuesBuilder.addAllStructValues(structObjs);
+      }
+    }
+
+    /**
+     * Create a struct for the given map based on object value type
+     *
+     * @param map The given map
+     * @return Struct The protobuf representation of the map using a struct object.
+     */
+    private static Struct createStruct(Map<String, Object> map) {
+      Builder struct = Struct.newBuilder();
+      for (Map.Entry<String, Object> entry : map.entrySet()) {
+        if (entry.getValue() instanceof Number) {
+          struct.putFields(
+              entry.getKey(),
+              Value.newBuilder().setNumberValue(((Number) entry.getValue()).doubleValue()).build());
+        } else if (entry.getValue() instanceof String) {
+          struct.putFields(
+              entry.getKey(), Value.newBuilder().setStringValue((String) entry.getValue()).build());
+        } else if (entry.getValue() instanceof Map) {
+          struct.putFields(
+              entry.getKey(),
+              Value.newBuilder()
+                  .setStructValue(createStruct((Map<String, Object>) entry.getValue()))
+                  .build());
+        }
+      }
+      return struct.build();
+    }
+
+    private static Struct createList(String[] objs) {
+      Builder struct = Struct.newBuilder();
+      for (Object obj : objs) {
+        struct.putFields("values", Value.newBuilder().setStringValue((String) obj).build());
+      }
+      return struct.build();
+    }
+
     /** Fetch field value from its doc value */
     private static void fetchFromDocVales(
         List<SearchResponse.Hit.Builder> sliceHits,

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/index/handlers/FieldUpdateHandler.java
@@ -37,7 +37,6 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import org.apache.lucene.queries.function.ValueSource;
 import org.apache.lucene.search.DoubleValuesSource;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/script/RuntimeScript.java
@@ -21,10 +21,7 @@
 import com.yelp.nrtsearch.server.luceneserver.script.RuntimeScript.SegmentFactory;
 import java.io.IOException;
 import java.util.Map;
-import java.util.Objects;
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.queries.function.FunctionValues;
-import org.apache.lucene.queries.function.ValueSource;
 
 /**
  * Script to produce a value for a given document. Implementations must have an execute function.
@@ -101,9 +98,7 @@ public interface Factory {
     SegmentFactory newFactory(Map<String, Object> params, DocLookup docLookup);
   }
 
-
   // compile context for the RuntimeScript, contains script type info
   public static final ScriptContext<Factory> CONTEXT =
-      new ScriptContext<>(
-          "runtime", Factory.class, SegmentFactory.class, RuntimeScript.class);
+      new ScriptContext<>("runtime", Factory.class, SegmentFactory.class, RuntimeScript.class);
 }
diff --git a/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java b/src/main/java/com/yelp/nrtsearch/server/luceneserver/search/SearchRequestProcessor.java
@@ -258,9 +258,9 @@ private static Map<String, FieldDef> getRuntimeFields(
       RuntimeScript.Factory factory =
           ScriptService.getInstance().compile(vf.getScript(), RuntimeScript.CONTEXT);
       Map<String, Object> params = ScriptParamsUtils.decodeParams(vf.getScript().getParamsMap());
-      RuntimeScript.SegmentFactory segmentFactory = factory.newFactory(params, indexState.docLookup);
-      FieldDef runtimeField =
-          new RuntimeFieldDef(vf.getName(), segmentFactory);
+      RuntimeScript.SegmentFactory segmentFactory =
+          factory.newFactory(params, indexState.docLookup);
+      FieldDef runtimeField = new RuntimeFieldDef(vf.getName(), segmentFactory);
       runtimeFields.put(vf.getName(), runtimeField);
     }
     return runtimeFields;
@@ -309,12 +309,12 @@ private static Map<String, FieldDef> getRetrieveFields(
    * @param queryFields mutable current map of query fields
    * @throws IllegalArgumentException if any index field already exists
    */
-  private static void addToQueryFields(Map<String, FieldDef> queryFields, Map<String, FieldDef> otherFields) {
+  private static void addToQueryFields(
+      Map<String, FieldDef> queryFields, Map<String, FieldDef> otherFields) {
     for (String key : otherFields.keySet()) {
       FieldDef current = queryFields.put(key, otherFields.get(key));
       if (current != null) {
-        throw new IllegalArgumentException(
-            "QueryFields: " + key + " specified multiple times");
+        throw new IllegalArgumentException("QueryFields: " + key + " specified multiple times");
       }
     }
   }