Run aggregate and group-by computes on the local engine

src/main/clojure/clojask/classes/DataFrame.clj

@@ -318,7 +318,9 @@
                                        (.indexOf index num))])
           aggre-func (mapv shift-func aggre-func)
           tmp (if (= ifheader true) (.printCol this output-dir select out))
-          res (start-onyx-aggre-only num-worker batch-size this output-dir exception aggre-func index select out)]
+          res (if (= (engine/engine) :local)
+                (engine/start-aggre-only num-worker batch-size this output-dir exception aggre-func index select out)
+                (start-onyx-aggre-only num-worker batch-size this output-dir exception aggre-func index select out))]
       (if (= res "success")
         "success"
         "failed")))
@@ -336,7 +338,9 @@
            data-index (mapv #(- % (count groupby-keys)) (remove #(< % (count groupby-keys)) select))
            groupby-index (vec (apply sorted-set (mapv #(nth % 1) (concat groupby-keys (u/gets aggre-keys data-index)))))
            mgroup (MGroup. (transient {}))
-           res (start-onyx-groupby num-worker batch-size this (if inmemory mgroup ".clojask/grouped/") groupby-keys groupby-index exception)]
+           local? (= (engine/engine) :local)
+           res ((if local? engine/start-groupby start-onyx-groupby)
+                num-worker batch-size this (if inmemory mgroup ".clojask/grouped/") groupby-keys groupby-index exception)]
        (if (= aggre-keys [])
          (println (str "Since the dataframe is only grouped by but not aggregated, the result will be the same as to choose the distinct values of "
                        "the groupby keys.")))
@@ -349,7 +353,8 @@
                formatter (.getFormatter (.col-info this))
                formatter (u/formatters-by-position formatter groupby-index)]
            (.final mgroup)
-           (if (= "success" (start-onyx-aggre num-worker batch-size this (if inmemory mgroup nil) output-dir exception aggre-func select formatter out))
+           (if (= "success" ((if local? engine/start-aggre start-onyx-aggre)
+                             num-worker batch-size this (if inmemory mgroup nil) output-dir exception aggre-func select formatter out))
              "success"
              (throw (OperationException. "Error when aggregating."))))
          (throw (OperationException. "Error when grouping by."))))

src/main/clojure/clojask/engine.clj

@@ -5,6 +5,8 @@
    environment variable CLOJASK_ENGINE, chooses the engine: \"onyx\" (the
    default) or \"local\". Drivers not ported yet always run on Onyx."
   (:require [clojure.java.io :as io]
+            [clojask.api.aggregate :refer [start]]
+            [clojask.groupby :as groupby]
             [clojask.utils :as u])
   (:import [com.clojask.exception ExecutionException]
            [java.util.concurrent ArrayBlockingQueue Callable ExecutorService
@@ -87,44 +89,151 @@
   [stage ^Throwable e]
   (ExecutionException. (format "[%s] (original error: %s)" stage (.getMessage e)) e))
 
-(defn- compute-worker
-  "The work on one batch of a plain compute: the selected columns of every
-   row that passes the filters, evaluated and formatted. Rows that fail a
-   filter are dropped. Same rows as onyx-comps/worker-func-gen-format."
-  [df exception index]
+(defn- row-worker
+  "The work on one batch of rows: the columns at index of every row that
+   passes the filters, evaluated and, if format? is true, formatted. Rows
+   that fail a filter are dropped. Same rows as onyx-comps/worker-func-gen
+   and worker-func-gen-format."
+  [df exception index format?]
   (let [col-info (:col-info df)
         operations (.getDesc col-info)
         types (.getType col-info)
         formats (.getFormatter col-info)
         filters (.getFilters (:row-info df))
         index (vec index)
-        eval-res (if exception u/eval-res u/eval-res-ne)]
+        eval-res (if exception u/eval-res u/eval-res-ne)
+        value (if format?
+                (fn [row i] ((or (get formats i) str) (eval-res row types formats operations i)))
+                (fn [row i] (eval-res row types formats operations i)))]
     (fn [batch]
       (into []
             (keep (fn [row]
                     (when (u/filter-check filters types row)
-                      (mapv (fn [i] ((or (get formats i) str) (eval-res row types formats operations i))) index))))
+                      (mapv (fn [i] (value row i)) index))))
             batch))))
 
+(defn- row-batches
+  "An open for run: the rows of df in batches of batch-size."
+  [df batch-size]
+  (fn []
+    (let [{:keys [data close]} ((.getFunc df) :map)]
+      {:data (partition-all batch-size data) :close close})))
+
+(defn- prepare
+  "(f), with a failure tagged as the preparing stage of stage."
+  [stage f]
+  (try
+    (f)
+    (catch Throwable e
+      (throw (stage-error (str "preparing stage (" stage ")") e)))))
+
+(defn- run-stage
+  "(f), with a failure tagged as the run stage of stage."
+  [stage f]
+  (try
+    (f)
+    (catch Throwable e
+      (throw (stage-error (str "run stage (" stage ")") e)))))
+
 (defn start-compute
   "The local engine's start-onyx: compute df into dist, appending after
    the header that DataFrame.compute may have written. Output is always in
    input order."
   [num-work batch-size df dist exception index melt out]
-  (let [work (try
-               (compute-worker df exception index)
-               (catch Throwable e
-                 (throw (stage-error "preparing stage (compute)" e))))]
-    (try
-      (with-open [wtr (io/writer dist :append true)]
-        (run num-work
-             (fn []
-               (let [{:keys [data close]} ((.getFunc df) :map)]
-                 {:data (partition-all batch-size data) :close close}))
-             work
-             (fn [rows]
-               (doseq [row rows]
-                 (out wtr (melt row))))))
-      (catch Throwable e
-        (throw (stage-error "run stage (compute)" e))))
+  (let [work (prepare "compute" #(row-worker df exception index true))]
+    (run-stage "compute"
+               #(with-open [wtr (io/writer dist :append true)]
+                  (run num-work (row-batches df batch-size) work
+                       (fn [rows]
+                         (doseq [row rows]
+                           (out wtr (melt row)))))))
+    "success"))
+
+(defn start-aggre-only
+  "The local engine's start-onyx-aggre-only: fold the rows that pass the
+   filters through aggre-func, pairs of a function and a position in
+   index, then append the result rows to dist. Rows are folded in input
+   order. Same output as clojask-aggre."
+  [num-work batch-size df dist exception aggre-func index select out]
+  (let [work (prepare "aggregate" #(row-worker df exception index false))
+        memo (volatile! (vec (repeat (count aggre-func) start)))]
+    (run-stage "aggregate"
+               (fn []
+                 (run num-work (row-batches df batch-size) work
+                      (fn [rows]
+                        (doseq [row rows]
+                          (vreset! memo (mapv (fn [[f i] prev] (f prev (nth row i))) aggre-func @memo)))))
+                 ;; a memo still holding the start sentinel saw no rows: write an empty cell
+                 (let [data (mapv (fn [m] (cond (identical? m start) [nil] (coll? m) m :else [m])) @memo)]
+                   (when-not (apply = (map count data))
+                     (throw (ExecutionException. "aggregation result is not of the same length")))
+                   (with-open [wtr (io/writer dist :append true)]
+                     (doseq [row (apply map vector data)]
+                       (out wtr [(u/gets row select)]))))))
+    "success"))
+
+(defn start-groupby
+  "The local engine's start-onyx-groupby: add every row that passes the
+   filters to its group, a file under dist when dist is a directory path,
+   otherwise the in-memory MGroup dist. Same groups as clojask-groupby."
+  [num-work batch-size df dist groupby-keys groupby-index exception & {:keys [format] :or {format false}}]
+  (let [col-info (:col-info df)
+        key-index (.getKeyIndex col-info)
+        formatter (.getFormatter col-info)
+        work (prepare "groupby" #(row-worker df exception (range (count key-index)) false))]
+    (run-stage "groupby"
+               #(run num-work (row-batches df batch-size) work
+                     (fn [rows]
+                       (doseq [row rows]
+                         (groupby/output-groupby dist row groupby-keys key-index formatter groupby-index format)))))
+    "success"))
+
+(defn- group-worker
+  "The work on one group of a group-aggregate, {:file f :d key}: the
+   group's result rows, the formatted key followed by the aggregates and
+   reordered by index. f is a group file, or the key in source when the
+   groups are in memory. Same rows as aggre-onyx-comps/worker-func-gen."
+  [df aggre-funcs index formatters source]
+  (let [groupby-index (mapv #(nth % 1) (.getGroupbyKeys (:row-info df)))
+        key-format (u/formatters-by-position (.getFormatter (:col-info df)) groupby-index)
+        key-index (range (count groupby-index))]
+    (fn [{:keys [file d]}]
+      (let [data (if (nil? source) (groupby/read-csv-seq file) (.getKey source file))
+            pre (u/gets-format d key-index key-format)
+            cols (zipmap (range) (apply map vector data))
+            res (reduce (fn [res [func i]]
+                          (let [new (func (get cols i))
+                                new (if (coll? new) new [new])
+                                new (mapv (fn [v] (if-let [formatter (get formatters i)] (formatter v) (str v))) new)]
+                            (if (or (= res []) (= (count new) (count (peek res))))
+                              (conj res new)
+                              (throw (Exception. "aggregation result is not of the same length")))))
+                        []
+                        aggre-funcs)]
+        (if (= res [])
+          [pre]
+          (mapv (fn [row] (u/gets (concat pre row) index)) (apply map vector res)))))))
+
+(defn- groups
+  "An open for run: one {:file f :d key} per group, from source when the
+   groups are in memory, else from the group files."
+  [source]
+  (fn []
+    {:data (if (nil? source)
+             (map (fn [file] {:file file :d (read-string (u/decode-str (.getName file)))})
+                  (rest (file-seq (io/file "./.clojask/grouped"))))
+             (map (fn [key] {:file key :d (read-string key)})
+                  (.getKeys source)))}))
+
+(defn start-aggre
+  "The local engine's start-onyx-aggre: aggregate every group on the pool
+   and append the result rows to dist, in the order the groups are listed.
+   Same output as aggre-output."
+  [num-work batch-size df source dist exception aggre-func index formatter out]
+  (let [work (prepare "groupby aggregate" #(group-worker df aggre-func index formatter source))]
+    (run-stage "groupby aggregate"
+               #(with-open [wtr (io/writer dist :append true)]
+                  (run num-work (groups source) work
+                       (fn [rows]
+                         (out wtr rows)))))
     "success"))

test/clojask/engine_test.clj

@@ -0,0 +1,106 @@
+(ns engine-test
+  "The same computes on the local engine and on Onyx. While the local
+   engine runs, starting Onyx fails the test. Starts Onyx."
+  (:require [clojure.test :refer :all]
+            [clojure.java.io :as io]
+            [clojask.dataframe :as ck]
+            [clojask.onyx-comps :as oc]
+            [clojask.api.aggregate :as agg]
+            [clojask.api.gb-aggregate :as gb])
+  (:import [com.clojask.exception ExecutionException]))
+
+(def input "test/clojask/Employees-example.csv")
+(def dir "test/clojask/test_outputs/engine/")
+
+(defn- on-engine
+  [engine f]
+  (let [old (System/getProperty "clojask.engine")]
+    (System/setProperty "clojask.engine" engine)
+    (try
+      (if (= engine "local")
+        (with-redefs [oc/with-onyx-env (fn [& _] (throw (AssertionError. "the local engine started Onyx")))]
+          (f))
+        (f))
+      (finally
+        (if old
+          (System/setProperty "clojask.engine" old)
+          (System/clearProperty "clojask.engine"))))))
+
+(defn- lines [path]
+  (with-open [r (io/reader path)] (vec (line-seq r))))
+
+(defn- result
+  "The header and the sorted rows that compute writes for the dataframe
+   (make-df) on engine, to a file or, with :in-memory, to the group store."
+  [engine make-df & opts]
+  (on-engine engine
+             (fn []
+               (let [path (str dir engine ".csv")]
+                 (io/make-parents path)
+                 (apply ck/compute (make-df) 2 path opts)
+                 (let [[header & rows] (lines path)]
+                   (into [header] (sort rows)))))))
+
+(defn- same-on-both
+  [make-df & opts]
+  (let [onyx (apply result "onyx" make-df opts)]
+    (is (< 1 (count onyx)) "the computed frame has rows")
+    (is (= onyx (apply result "local" make-df opts)))))
+
+(defn- salaries []
+  (doto (ck/dataframe input)
+    (ck/set-type "Salary" "double")))
+
+(deftest aggregate-only
+  (same-on-both #(doto (salaries)
+                   (ck/filter "Salary" (fn [s] (<= s 800)))
+                   (ck/aggregate agg/max ["Salary"] ["max"])
+                   (ck/aggregate agg/min ["Salary"] ["min"])
+                   (ck/aggregate agg/count ["Employee"] ["n"]))))
+
+(deftest group-aggregate
+  (doseq [in-memory [false true]]
+    (testing (str "in-memory " in-memory)
+      (same-on-both #(doto (salaries)
+                       (ck/group-by ["Department"])
+                       (ck/aggregate gb/max ["Salary"] ["max"])
+                       (ck/aggregate gb/mean ["Salary"] ["mean"]))
+                    :in-memory in-memory))))
+
+(deftest group-aggregate-multi-row-result
+  (same-on-both #(doto (salaries)
+                   (ck/group-by ["Department"])
+                   (ck/aggregate gb/smallest3 ["Salary"] ["low"]))
+                :select ["low" "Department"]))
+
+(deftest group-by-only
+  (same-on-both #(doto (ck/dataframe input)
+                   (ck/group-by ["Department"]))))
+
+(def ^:private boom?
+  "Set once the dataframe is built, so the preview that checks each
+   aggregate does not throw."
+  (atom false))
+
+(defn- boom [f]
+  (fn [& args]
+    (if @boom?
+      (throw (RuntimeException. "boom"))
+      (apply f args))))
+
+(deftest aggregate-failure-is-staged
+  (doseq [engine ["onyx" "local"]
+          [label make-df] [["aggregate" #(ck/aggregate (salaries) (boom agg/max) ["Salary"] ["x"])]
+                           ["group aggregate" #(doto (salaries)
+                                                 (ck/group-by ["Department"])
+                                                 (ck/aggregate (boom gb/max) ["Salary"] ["x"]))]]]
+    (testing (str engine " " label)
+      (let [df (make-df)]
+        (reset! boom? true)
+        (try
+          (is (thrown-with-msg? ExecutionException #"boom"
+                                (on-engine engine #(ck/compute df 2 (str dir "failed.csv")))))
+          (finally
+            (reset! boom? false))))
+      (is (empty? (rest (file-seq (io/file ".clojask/grouped"))))
+          "the group files are deleted"))))