From 10cfb08a38117e708e7c114a9b732978311fa4ae Mon Sep 17 00:00:00 2001
From: Joel Bernstein <jbernste@apache.org>
Date: Tue, 10 Dec 2019 12:57:05 -0500
Subject: [PATCH] SOLR-14043: Allow the precision Stream Evaluator to operate
 on matrices

---
 .../solrj/io/eval/PrecisionEvaluator.java     | 13 +++++++++-
 .../solrj/io/stream/MathExpressionTest.java   | 26 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/PrecisionEvaluator.java b/solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/PrecisionEvaluator.java
index 532f2ea40c4..09df0837648 100644
--- a/solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/PrecisionEvaluator.java
+++ b/solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/PrecisionEvaluator.java
@@ -25,7 +25,7 @@ import org.apache.commons.math3.util.Precision;
 import org.apache.solr.client.solrj.io.stream.expr.StreamExpression;
 import org.apache.solr.client.solrj.io.stream.expr.StreamFactory;
 
-public class PrecisionEvaluator extends RecursiveNumericEvaluator implements TwoValueWorker {
+public class PrecisionEvaluator extends RecursiveObjectEvaluator implements TwoValueWorker {
   protected static final long serialVersionUID = 1L;
 
   public PrecisionEvaluator(StreamExpression expression, StreamFactory factory) throws IOException{
@@ -43,6 +43,17 @@ public class PrecisionEvaluator extends RecursiveNumericEvaluator implements Two
     }
     else if(value instanceof List){
       return ((List<?>)value).stream().map(innerValue -> doWork(innerValue, ((Number)value2).intValue())).collect(Collectors.toList());
+    } else if(value instanceof Matrix) {
+      int p = ((Number)value2).intValue();
+      Matrix matrix = (Matrix)value;
+      double[][] data = matrix.getData();
+      for(int i=0; i<data.length; ++i) {
+        for(int j=0; j < data[i].length; j++) {
+          data[i][j] = Precision.round(data[i][j], p);
+        }
+      }
+
+      return matrix;
     }
     else{
       return Precision.round(((Number)value).doubleValue(), ((Number)value2).intValue());
diff --git a/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/MathExpressionTest.java b/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/MathExpressionTest.java
index ae7eb689bf2..25de9e1b194 100644
--- a/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/MathExpressionTest.java
+++ b/solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/MathExpressionTest.java
@@ -5463,6 +5463,32 @@ public class MathExpressionTest extends SolrCloudTestCase {
     assertEquals(num, 1.4445, 0.0);
   }
 
+  @Test
+  public void testPrecisionMatrix() throws Exception {
+    String cexpr = "let(a=matrix(array(1.3333999, 2.4444445), array(2.333333, 10.10009)), b=precision(a, 4))";
+    ModifiableSolrParams paramsLoc = new ModifiableSolrParams();
+    paramsLoc.set("expr", cexpr);
+    paramsLoc.set("qt", "/stream");
+    String url = cluster.getJettySolrRunners().get(0).getBaseUrl().toString()+"/"+COLLECTIONORALIAS;
+    TupleStream solrStream = new SolrStream(url, paramsLoc);
+    StreamContext context = new StreamContext();
+    solrStream.setStreamContext(context);
+    List<Tuple> tuples = getTuples(solrStream);
+    assertTrue(tuples.size() == 1);
+
+    List<List<Number>> rows = (List<List<Number>>)tuples.get(0).get("b");
+    assertTrue(rows.size() == 2);
+    List<Number> row1 = rows.get(0);
+    assertTrue(row1.size() == 2);
+    assertEquals(row1.get(0).doubleValue(), 1.3334, 0);
+    assertEquals(row1.get(1).doubleValue(),  2.4444, 0);
+
+    List<Number> row2 = rows.get(1);
+    assertTrue(row2.size() == 2);
+    assertEquals(row2.get(0).doubleValue(), 2.3333, 0);
+    assertEquals(row2.get(1).doubleValue(),  10.1001, 0);
+  }
+
   @Test
   public void testMinMaxScale() throws Exception {
     String cexpr = "let(echo=true, a=minMaxScale(matrix(array(1,2,3,4,5), array(10,20,30,40,50))), " +