From 4295d9d654c7592c702cc894c90dd396a20bdaeb Mon Sep 17 00:00:00 2001 From: Matt Buehlmaier Date: Tue, 13 Oct 2026 16:50:18 +0800 Subject: [PATCH] Run aggregate and group-by computes on the local engine With clojask.engine=local, aggregate-only, group-by and group-aggregate computes no longer start Onyx. engine/start-aggre-only folds the rows in input order on the calling thread, engine/start-groupby writes the rows to their group files or the in-memory groups, and engine/start-aggre aggregates one group per pool task. They take the same arguments as the Onyx drivers they replace. Joins still run on Onyx, including their group-by phase. engine-test runs each of these computes on both engines and compares the sorted output, with Onyx made unreachable during the local run. It also checks that a failing aggregate surfaces as an ExecutionException on both engines and leaves no group files. --- .../clojure/clojask/classes/DataFrame.clj | 11 +- src/main/clojure/clojask/engine.clj | 155 +++++++++++++++--- test/clojask/engine_test.clj | 106 ++++++++++++ 3 files changed, 246 insertions(+), 26 deletions(-) create mode 100644 test/clojask/engine_test.clj diff --git a/src/main/clojure/clojask/classes/DataFrame.clj b/src/main/clojure/clojask/classes/DataFrame.clj index b552979..2aafabb 100644 --- a/src/main/clojure/clojask/classes/DataFrame.clj +++ b/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.")))) diff --git a/src/main/clojure/clojask/engine.clj b/src/main/clojure/clojask/engine.clj index c5cc5e1..ecf0458 100644 --- a/src/main/clojure/clojask/engine.clj +++ b/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")) diff --git a/test/clojask/engine_test.clj b/test/clojask/engine_test.clj new file mode 100644 index 0000000..84d1527 --- /dev/null +++ b/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"))))