SOLR-11602: Add Markov Chain Stream Evaluator

solr/core/src/java/org/apache/solr/handler/StreamHandler.java

@@ -274,6 +274,7 @@ public class StreamHandler extends RequestHandlerBase implements SolrCoreAware,
         .withFunctionName("triangularDistribution", TriangularDistributionEvaluator.class)
         .withFunctionName("precision", PrecisionEvaluator.class)
         .withFunctionName("minMaxScale", MinMaxScaleEvaluator.class)
+        .withFunctionName("markovChain", MarkovChainEvaluator.class)
 
         // Boolean Stream Evaluators
 

solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/MarkovChainEvaluator.java

@@ -0,0 +1,102 @@
+/*
+ * 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.eval;
+
+import java.io.IOException;
+
+import java.util.Locale;
+import java.util.Random;
+import java.util.List;
+import java.util.ArrayList;
+
+import org.apache.commons.math3.distribution.EnumeratedIntegerDistribution;
+import org.apache.commons.math3.util.MathArrays;
+import org.apache.solr.client.solrj.io.stream.expr.StreamExpression;
+import org.apache.solr.client.solrj.io.stream.expr.StreamFactory;
+
+public class MarkovChainEvaluator extends RecursiveObjectEvaluator implements ManyValueWorker {
+  protected static final long serialVersionUID = 1L;
+
+  public MarkovChainEvaluator(StreamExpression expression, StreamFactory factory) throws IOException{
+    super(expression, factory);
+
+    if(2 < containedEvaluators.size()){
+      throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - expecting no more then two parameters but found %d",expression,containedEvaluators.size()));
+    }
+  }
+
+  @Override
+  public Object doWork(Object... values) throws IOException{
+
+    int state = -1;
+
+    if(values.length == 2) {
+      state = ((Number)values[1]).intValue();
+    }
+
+    if(values[0] instanceof Matrix) {
+      Matrix matrix = (Matrix) values[0];
+      return new MarkovChain(matrix, state);
+    } else {
+      throw new IOException("matrix parameter expected for transpose function");
+    }
+  }
+
+  public static class MarkovChain {
+
+    private int state;
+    private EnumeratedIntegerDistribution[] distributions;
+
+    public MarkovChain(Matrix matrix, int state) throws IOException {
+      double[][] data = matrix.getData();
+
+      if(data.length != data[0].length) {
+        throw new IOException("Markov chain must be initialized with a square matrix.");
+      }
+
+      this.distributions = new EnumeratedIntegerDistribution[data.length];
+
+      if(state > -1) {
+        this.state = state;
+      } else {
+        this.state = new Random().nextInt(data.length);
+      }
+
+      for(int i=0; i<data.length; i++) {
+        double[] probabilities = data[i];
+
+        //Create the states array needed by the enumerated distribution
+        int[] states = MathArrays.sequence(data.length, 0, 1);
+        distributions[i] = new EnumeratedIntegerDistribution(states, probabilities);
+      }
+    }
+
+    public Number sample() {
+      this.state = distributions[this.state].sample();
+      return this.state;
+    }
+
+    public int[] sample(int size) {
+      int[] sample = new int[size];
+      for(int i=0; i<size; i++) {
+        sample[i] = sample().intValue();
+      }
+
+      return sample;
+    }
+  }
+}

solr/solrj/src/java/org/apache/solr/client/solrj/io/eval/SampleEvaluator.java

@@ -43,8 +43,8 @@ public class SampleEvaluator extends RecursiveObjectEvaluator implements ManyVal
 
     Object first = objects[0];
 
-    if(!(first instanceof RealDistribution) && !(first instanceof IntegerDistribution)){
-      throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - found type %s for the first value, expecting a Real or Integer Distribution",toExpression(constructingFactory), first.getClass().getSimpleName()));
+    if(!(first instanceof RealDistribution) && !(first instanceof IntegerDistribution) && !(first instanceof MarkovChainEvaluator.MarkovChain)){
+      throw new IOException(String.format(Locale.ROOT,"Invalid expression %s - found type %s for the first value, expecting a Markov Chain, Real or Integer Distribution",toExpression(constructingFactory), first.getClass().getSimpleName()));
     }
 
     Object second = null;
@@ -52,7 +52,14 @@ public class SampleEvaluator extends RecursiveObjectEvaluator implements ManyVal
       second = objects[1];
     }
 
-    if(first instanceof RealDistribution) {
+    if(first instanceof MarkovChainEvaluator.MarkovChain) {
+      MarkovChainEvaluator.MarkovChain markovChain = (MarkovChainEvaluator.MarkovChain)first;
+      if(second != null) {
+        return Arrays.stream(markovChain.sample(((Number) second).intValue())).mapToObj(item -> item).collect(Collectors.toList());
+      } else {
+        return markovChain.sample();
+      }
+    } else if (first instanceof RealDistribution) {
       RealDistribution realDistribution = (RealDistribution) first;
       if(second != null) {
         return Arrays.stream(realDistribution.sample(((Number) second).intValue())).mapToObj(item -> item).collect(Collectors.toList());

solr/solrj/src/test/org/apache/solr/client/solrj/io/stream/StreamExpressionTest.java

@@ -6162,6 +6162,34 @@ public class StreamExpressionTest extends SolrCloudTestCase {
     assertEquals(array2.get(2).doubleValue(), 1, 0.0);
   }
 
+  @Test
+  public void testMarkovChain() throws Exception {
+    String cexpr = "let(state0=array(.5,.5),\n" +
+                   "    state1=array(.5,.5),\n" +
+                   "    states=matrix(state0, state1),\n" +
+                   "    m=markovChain(states, 0),\n" +
+                   "    s=sample(m, 50000),\n" +
+                   "    f=freqTable(s))";
+    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<Map<String, Number>> out = (List<Map<String, Number>>)tuples.get(0).get("f");
+    assertEquals(out.size(), 2);
+    Map<String, Number> bin0 = out.get(0);
+    double state0Pct = bin0.get("pct").doubleValue();
+    assertEquals(state0Pct, .5, .015);
+    Map<String, Number> bin1 = out.get(1);
+    double state1Pct = bin1.get("pct").doubleValue();
+    assertEquals(state1Pct, .5, .015);
+  }
+
+
 
   @Test
   public void testAddAll() throws Exception {