From 07852efd473c9ed21bdc79e8d468a14347baf1e7 Mon Sep 17 00:00:00 2001
From: Anand Avati <avati@redhat.com>
Date: Tue, 23 Jun 2015 20:23:53 -0700
Subject: [PATCH] MAHOUT-1736: Implement allreduceBlock() on H2O

Signed-off-by: Anand Avati <avati@redhat.com>
---
 .../apache/mahout/h2obindings/H2OHelper.java  | 32 +++++++++++++++++++
 .../apache/mahout/h2obindings/H2OEngine.scala |  2 +-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/h2o/src/main/java/org/apache/mahout/h2obindings/H2OHelper.java b/h2o/src/main/java/org/apache/mahout/h2obindings/H2OHelper.java
index 0fae5a8148..c9d91f9009 100644
--- a/h2o/src/main/java/org/apache/mahout/h2obindings/H2OHelper.java
+++ b/h2o/src/main/java/org/apache/mahout/h2obindings/H2OHelper.java
@@ -33,14 +33,19 @@
 
 import java.util.Map;
 import java.util.HashMap;
+import java.io.Serializable;
 
 import org.apache.mahout.h2obindings.drm.H2ODrm;
+import org.apache.mahout.h2obindings.drm.H2OBCast;
 
 // for makeEmptyStrVec
 import water.Key;
 import water.DKV;
 import water.fvec.CStrChunk;
 
+import scala.Function1;
+import scala.Function2;
+
 /**
  * Collection of helper methods for H2O backend.
  */
@@ -437,4 +442,31 @@ public static int vecChunkLen(Vec template, int chunk) {
   public static H2ODrm emptyDrm(long nrow, int ncol, int minHint, int exactHint) {
     return new H2ODrm(emptyFrame(nrow, ncol, minHint, exactHint));
   }
+
+  public static Matrix allreduceBlock(H2ODrm drmA, Object bmfn, Object rfn) {
+    class MRTaskMR extends MRTask<MRTaskMR> {
+      H2OBCast<Matrix> bmf_out;
+      Serializable bmf;
+      Serializable rf;
+
+      public MRTaskMR(Object _bmf, Object _rf) {
+        bmf = (Serializable) _bmf;
+        rf = (Serializable) _rf;
+      }
+
+      @Override
+      public void map(Chunk chks[]) {
+        Function1 f = (Function1) bmf;
+        bmf_out = new H2OBCast((Matrix)f.apply(new scala.Tuple2(null, new H2OBlockMatrix(chks))));
+      }
+
+      @Override
+      public void reduce(MRTaskMR that) {
+        Function2 f = (Function2) rf;
+        bmf_out = new H2OBCast((Matrix)f.apply(this.bmf_out.value(), that.bmf_out.value()));
+      }
+    }
+
+    return new MRTaskMR(bmfn, rfn).doAll(drmA.frame).bmf_out.value();
+  }
 }
diff --git a/h2o/src/main/scala/org/apache/mahout/h2obindings/H2OEngine.scala b/h2o/src/main/scala/org/apache/mahout/h2obindings/H2OEngine.scala
index 420c22d82a..e6569f9832 100644
--- a/h2o/src/main/scala/org/apache/mahout/h2obindings/H2OEngine.scala
+++ b/h2o/src/main/scala/org/apache/mahout/h2obindings/H2OEngine.scala
@@ -167,7 +167,7 @@ object H2OEngine extends DistributedEngine {
    *
    */
   override def allreduceBlock[K: ClassTag](drm: CheckpointedDrm[K], bmf: BlockMapFunc2[K], rf: BlockReduceFunc)
-  : Matrix = ???
+  : Matrix = H2OHelper.allreduceBlock(drm.h2odrm, bmf, rf)
 
   /**
    * TODO: implement this please.