SOLR-16287 : MapStream - remap tuple value(s)

solr/solr-ref-guide/modules/query-guide/pages/stream-decorator-reference.adoc

@@ -1009,6 +1009,62 @@ list(tuple(a=search(collection1, q="*:*", fl="id, prod_ss", sort="id asc")),
      tuple(a=search(collection2, q="*:*", fl="id, prod_ss", sort="id asc")))
 ----
 
+[#map_expression]
+== map
+
+Given any number of "key=newKey" mappings, the `map` function wraps a Stream Expression
+and re-maps 'key=value' tuples into 'key=tuple(newKey,value)'
+
+=== map Parameters
+
+* `incoming stream`: (Mandatory) A single incoming stream.
+* `"key=newKey"` {1,*} : re-map tuples with `key=value` to `key=tuple(newKey,value)`
+
+=== map Syntax
+
+[source,text]
+----
+Given:
+
+{"id": "1",
+ "count_ii": "2"
+}
+
+Then:
+map(tuple(id=1,count_ii=2),"count_ii=add-distinct")
+
+Will output:
+{
+"count_ii": {
+  "add-distinct": "2"
+ },
+ "id": "1"
+}
+
+e.g. atomically update some records
+
+update( destinationCollection,
+        batchSize=500,
+        map(
+          search(collection1,
+                 q=*:*,
+                 qt="/export",
+                 fl="id,a_i",
+                 sort="a_i asc")),
+        "a_i=set"
+)
+
+----
+
+Expression above sends the tuples returned by the `search` function to the `destinationCollection` to be atomically
+updated, the rest of the document will remain as it was, existing requirements for atomic updates still apply.
+Assumes the id in `collection1` and `destinationCollection` represent the same thing.
+
+Here we atomically update `a_i` by setting it to values from the `collection1` collection for
+`collection1.id == destinationCollection.id`
+
+Without the map decorator, the update would simply overwrite any existing document.
+
 == merge
 
 The `merge` function merges two or more streaming expressions and maintains the ordering of the underlying streams.

solr/solrj/src/java/org/apache/solr/client/solrj/io/Lang.java

@@ -57,6 +57,7 @@ public static void register(StreamFactory streamFactory) {
 
         // decorator streams
         .withFunctionName("merge", MergeStream.class)
+        .withFunctionName("map", MapStream.class)
         .withFunctionName("unique", UniqueStream.class)
         .withFunctionName("top", RankStream.class)
         .withFunctionName("group", GroupOperation.class)

solr/solrj/src/java/org/apache/solr/client/solrj/io/stream/MapStream.java

@@ -0,0 +1,151 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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.
+ */
+
+package org.apache.solr.client.solrj.io.stream;
+
+import java.io.IOException;
+import java.util.*;
+import org.apache.solr.client.solrj.io.Tuple;
+import org.apache.solr.client.solrj.io.comp.StreamComparator;
+import org.apache.solr.client.solrj.io.stream.expr.*;
+
+public class MapStream extends TupleStream implements Expressible {
+
+  private static final long serialVersionUID = 1;
+
+  private final TupleStream stream;
+
+  private final Map<String, String> mappings;
+
+  public MapStream(TupleStream stream, Map<String, String> mappings) {
+    this.stream = stream;
+    this.mappings = mappings;
+  }
+
+  public MapStream(StreamExpression expression, StreamFactory factory) throws IOException {
+    List<StreamExpression> streamExpressions =
+        factory.getExpressionOperandsRepresentingTypes(
+            expression, Expressible.class, TupleStream.class);
+    if (1 != streamExpressions.size()) {
+      throw new IOException(
+          String.format(
+              Locale.ROOT,
+              "Invalid expression %s - expecting single stream but found %d (must be TupleStream types)",
+              expression,
+              streamExpressions.size()));
+    }
+
+    stream = factory.constructStream(streamExpressions.get(0));
+
+    List<StreamExpressionParameter> mapFieldExpressions =
+        factory.getOperandsOfType(expression, StreamExpressionValue.class);
+    if (0 == mapFieldExpressions.size()) {
+      throw new IOException(
+          String.format(
+              Locale.ROOT,
+              "Invalid expression %s - expecting at least one select field but found %d",
+              expression,
+              streamExpressions.size()));
+    }
+
+    mappings = new HashMap<>();
+    mapFieldExpressions.forEach(
+        op -> {
+          String value = ((StreamExpressionValue) op).getValue().trim();
+          if (value.length() > 2 && value.startsWith("\"") && value.endsWith("\"")) {
+            value = value.substring(1, value.length() - 1);
+          }
+          String[] parts = value.split("=");
+          mappings.put(parts[0].trim(), parts[1].trim());
+        });
+  }
+
+  @Override
+  public void setStreamContext(StreamContext context) {
+    this.stream.setStreamContext(context);
+  }
+
+  @Override
+  public List<TupleStream> children() {
+    List<TupleStream> l = new ArrayList<>();
+    l.add(stream);
+    return l;
+  }
+
+  @Override
+  public void open() throws IOException {
+    stream.open();
+  }
+
+  @Override
+  public void close() throws IOException {
+    stream.close();
+  }
+
+  @Override
+  public Tuple read() throws IOException {
+    Tuple original = stream.read();
+    if (original.EOF) return original;
+    final Tuple workingToReturn = new Tuple();
+    original
+        .getFields()
+        .forEach(
+            (k, v) ->
+                workingToReturn.put(
+                    k, mappings.containsKey(k) ? new Tuple(mappings.get(k), v) : v));
+    return workingToReturn;
+  }
+
+  @Override
+  public StreamComparator getStreamSort() {
+    return null;
+  }
+
+  @Override
+  public StreamExpressionParameter toExpression(StreamFactory factory) throws IOException {
+    return toExpression(factory, true);
+  }
+
+  private StreamExpression toExpression(StreamFactory factory, boolean includeStreams)
+      throws IOException {
+    StreamExpression expression = new StreamExpression(factory.getFunctionName(this.getClass()));
+    if (includeStreams) {
+      // stream
+      if (stream instanceof Expressible) {
+        expression.addParameter(((Expressible) stream).toExpression(factory));
+      } else {
+        throw new IOException(
+            "This SelectStream contains a non-expressible TupleStream - it cannot be converted to an expression");
+      }
+    } else {
+      expression.addParameter("<stream>");
+    }
+    mappings.forEach((k, v) -> expression.addParameter(k + "=" + v));
+
+    return expression;
+  }
+
+  @Override
+  public Explanation toExplanation(StreamFactory factory) throws IOException {
+    return new StreamExplanation(getStreamNodeId().toString())
+        .withChildren(new Explanation[] {stream.toExplanation(factory)})
+        .withFunctionName(factory.getFunctionName(this.getClass()))
+        .withImplementingClass(this.getClass().getName())
+        .withExpressionType(Explanation.ExpressionType.STREAM_DECORATOR)
+        .withExpression(toExpression(factory, false).toString());
+  }
+}

solr/solrj/src/test/org/apache/solr/client/solrj/io/TestLang.java

@@ -117,6 +117,7 @@ public class TestLang extends SolrTestCase {
     "copyOfRange",
     "copyOf",
     "cov",
+    "map",
     "corr",
     "describe",
     "distance",