Performance benchmark data (#134)

* Added rand-int and rand-uniform

* Added random-int and random-unif

* Added random-exp + clean up

* Bug fixes for random number generation + partition-by for writer

* Cleaner data generation script

* Fixed lint error

* Added codecov yml

* Added public token

* Take down code-cov token

* Renamed codecov

* Added public token

* Moved codecov token

Author: anthony-khong

.github/workflows/continuous-integration.yml

@@ -28,6 +28,8 @@ jobs:
           entrypoint: lein
           args: coverage
       - uses: codecov/codecov-action@v1
+        env:
+            CODECOV_TOKEN: bb12d419-f860-48ff-a3bf-8fca4114cc3d
         with:
          file: target/coverage/codecov.json
          fail_ci_if_error: true

codecov.yml

@@ -0,0 +1,20 @@
+codecov:
+  require_ci_to_pass: yes
+
+coverage:
+  precision: 3
+  round: down
+  range: "70...100"
+
+parsers:
+  gcov:
+    branch_detection:
+      conditional: yes
+      loop: yes
+      method: no
+      macro: no
+
+comment:
+  layout: "reach,diff,flags,tree"
+  behavior: default
+  require_changes: no

examples/performance_benchmark_data.clj

@@ -0,0 +1,58 @@
+(ns examples.performance-benchmark-data
+  (:require
+    [zero-one.geni.core :as g]))
+
+(def data-path "/data/performance-benchmark-data")
+
+(def skeleton-df
+  (g/cache (g/table->dataset (repeat (int 2e6) [1]) [:dummy])))
+
+(defn transaction-id-col []
+  (g/concat (g/str (g/random-int))
+            (g/lit "-")
+            (g/str (g/random-int))
+            (g/lit "-")
+            (g/str (g/random-int))))
+
+(def date-col
+  (g/concat :year (g/lit "-") :month (g/lit "-") :day))
+
+(def max-days {1  31
+               2  28
+               3  31
+               4  30
+               5  31
+               6  30
+               7  31
+               8  31
+               9  30
+               10 31
+               11 30
+               12 30})
+
+(doall
+  (for [month (range 1 13)]
+    (-> skeleton-df
+        (g/select
+          {:trx-id    (transaction-id-col)
+           :member-id (g/int (g/rexp 5e-6))
+           :quantity  (g/int (g/inc (g/rexp)))
+           :price     (g/pow 2 (g/random-int 16 20))
+           :style-id  (g/int (g/rexp 1e-2))
+           :brand-id  (g/int (g/rexp 1e-2))
+           :year      2019
+           :month     month
+           :day       (g/random-int 1 (inc (max-days month)))})
+        (g/with-column :date (g/to-date date-col))
+        (g/coalesce 1)
+        (g/write-parquet! data-path {:mode "append"}))))
+
+(comment
+
+  (-> (g/read-parquet! data-path)
+      (g/group-by :trx-id)
+      g/count
+      (g/describe :count)
+      g/show)
+
+  true)

src/zero_one/geni/core.clj

@@ -49,6 +49,7 @@
                             partition-by
                             pos?
                             rand
+                            rand-int
                             rand-nth
                             remove
                             rename-keys
@@ -149,6 +150,7 @@
 (import-vars
   [zero-one.geni.sql
    !
+   **
    ->date-col
    ->timestamp-col
    ->utc-timestamp
@@ -565,6 +567,15 @@
    nunique
    qcut
    random-choice
+   random-exp
+   random-int
+   random-norm
+   random-uniform
+   rchoice
+   rexp
+   rnorm
+   runif
+   runiform
    select-columns
    shape
    value-counts])

src/zero_one/geni/data_sources.clj

@@ -2,7 +2,8 @@
   (:refer-clojure :exclude [partition-by sort-by])
   (:require
     [camel-snake-kebab.core :refer [->camelCase]]
-    [zero-one.geni.defaults])
+    [zero-one.geni.defaults]
+    [zero-one.geni.utils :refer [ensure-coll]])
   (:import
     (org.apache.spark.sql SparkSession)))
 
@@ -82,15 +83,21 @@
          configured-reader   (configure-reader-or-writer unconfigured-reader options)]
      (.load configured-reader))))
 
+(defn- partition-by-arg [partition-id]
+  (into-array java.lang.String (map name (ensure-coll partition-id))))
+
 (defn write-data! [format dataframe path options]
   (let [mode                (:mode options)
+        partition-id        (:partition-by options)
         unconfigured-writer (-> dataframe
                                 (.write)
                                 (.format format)
-                                (cond-> mode (.mode mode)))
+                                (cond-> mode (.mode mode))
+                                (cond-> partition-id
+                                  (.partitionBy (partition-by-arg partition-id))))
         configured-writer   (configure-reader-or-writer
                               unconfigured-writer
-                              (dissoc options :mode))]
+                              (dissoc options :mode :partition-by))]
     (.save configured-writer path)))
 
 (defn write-parquet!

src/zero_one/geni/foreign_idioms.clj

@@ -23,17 +23,52 @@
                                 (str low)
                                 (str high))))))
 
+(defn random-uniform
+  ([] (random-uniform 0.0 1.0))
+  ([low high] (random-uniform low high (rand-int Integer/MAX_VALUE)))
+  ([low high seed]
+   (let [length (Math/abs (- high low))
+         base   (min high low)]
+     (column/+ base (column/* length (sql/rand seed))))))
+(def runiform random-uniform)
+(def runif random-uniform)
+
+(defn random-norm
+  ([] (random-norm 0.0 1.0))
+  ([mu sigma] (random-norm mu sigma (rand-int Integer/MAX_VALUE)))
+  ([mu sigma seed] (column/+ mu (column/* sigma (sql/randn seed)))))
+(def rnorm random-norm)
+
+(defn random-exp
+  ([] (random-exp 1.0))
+  ([rate] (random-exp rate (rand-int Integer/MAX_VALUE)))
+  ([rate seed] (-> (sql/rand seed)
+                   sql/log
+                   (column/* -1.0)
+                   (column// rate))))
+(def rexp random-exp)
+
+(defn random-int
+  ([] (random-int 0 (dec Integer/MAX_VALUE)))
+  ([low high] (random-int low high (rand-int Integer/MAX_VALUE)))
+  ([low high seed]
+   (let [length (Math/abs (- high low))
+         base   (min high low)
+         ->long #(column/cast % "long")]
+     (column/+ (->long base) (->long (column/* length (sql/rand seed)))))))
+
 (defn random-choice
   ([choices]
    (let [n-choices (count choices)]
      (random-choice choices (take n-choices (repeat (/ 1.0 n-choices))))))
-  ([choices probs]
+  ([choices probs] (random-choice choices probs (rand-int Integer/MAX_VALUE)))
+  ([choices probs seed]
    (assert (and (= (count choices) (count probs))
                 (every? pos? probs))
            "random-choice args must have same lengths.")
    (assert (< (Math/abs (- (apply + probs) 1.0)) 1e-4)
            "random-choice probs must some to one.")
-   (let [rand-col    (column/->column (sql/rand))
+   (let [rand-col    (column/->column (sql/rand seed))
          cum-probs   (reductions + probs)
          choice-cols (map (fn [choice prob]
                             (sql/when (column/< rand-col (+ prob 1e-6))
@@ -42,6 +77,7 @@
                           cum-probs)]
      (.as (apply polymorphic/coalesce choice-cols)
           (format "choice(%s, %s)" (str choices) (str probs))))))
+(def rchoice random-choice)
 
 ;; Pandas
 (defn value-counts [dataframe]

src/zero_one/geni/main.clj

@@ -32,12 +32,12 @@
 
   (require '[clojure.reflect :as r])
   (import '(org.apache.spark.sql Dataset))
-  (->> (r/reflect Dataset)
+  (->> (r/reflect (.write dataframe))
        :members
-       (clojure.core/filter #(= (:name %) 'toDF))
+       (clojure.core/filter #(= (:name %) 'partitionBy))
        ;(mapv :parameter-types)
        ;(clojure.core/filter #(= (:name %) 'toDF))
        ;clojure.core/sort
-       pprint)
+       clojure.pprint/pprint)
 
   0)

src/zero_one/geni/sql.clj

@@ -213,6 +213,7 @@
 (defn log2 [expr] (functions/log2 (->column expr)))
 (defn pmod [left-expr right-expr] (functions/pmod (->column left-expr) (->column right-expr)))
 (defn pow [base exponent] (functions/pow (->column base) (->column exponent)))
+(def ** pow)
 (defn radians [expr] (functions/radians (->column expr)))
 (defn rint [expr] (functions/rint (->column expr)))
 (defn round [expr] (functions/round (->column expr)))

test/zero_one/geni/data_sources_test.clj

@@ -118,3 +118,12 @@
                                    :url     (str "jdbc:sqlite:" temp-file)
                                    :dbtable "housing"}))]
     (g/collect-vals write-df) => (g/collect-vals read-df)))
+
+(fact "Can write parquet with :partition-by option" :slow
+  (let [temp-file (.toString (create-temp-file! ".parquet"))
+        read-df  (do (g/write-parquet!
+                       write-df
+                       temp-file
+                       {:mode "overwrite" :partition-by [:Method]})
+                     (g/read-parquet! temp-file))]
+    (set (g/collect write-df)) => (set (g/collect read-df))))

test/zero_one/geni/numpy_test.clj

@@ -4,6 +4,51 @@
     [zero-one.geni.core :as g]
     [zero-one.geni.test-resources :refer [df-20]]))
 
+(defn descriptive-stats [col]
+  (-> (g/table->dataset (mapv vector (range 200)) [:idx])
+      (g/with-column :x col)
+      (g/agg {:min  (g/min :x)
+              :mean (g/mean :x)
+              :std  (g/stddev :x)
+              :max  (g/max :x)})
+      g/first))
+
+(fact "On random-exp" :slow
+  (descriptive-stats (g/random-exp)) => #(and (< 0.5 (:mean %) 1.5)
+                                              (< 0.5 (:std %) 1.5))
+  (descriptive-stats (g/random-exp 5)) => #(and (< 0.1 (:mean %) 0.3)
+                                                (< 0.1 (:std %) 0.3)))
+
+(fact "On random-norm" :slow
+  (descriptive-stats (g/random-norm)) => #(and (< -0.3 (:mean %) 0.3)
+                                               (< -0.8 (:std %) 1.2))
+  (descriptive-stats (g/random-norm -3 2)) => #(and (< -4.0 (:mean %) -2.0)
+                                                    (< 1.5 (:std %) 2.5)))
+
+(fact "On random-int" :slow
+  (descriptive-stats (g/random-int)) => #(and (pos? (:max %))
+                                              (pos? (:min %))
+                                              (integer? (:max %))
+                                              (integer? (:min %)))
+  (descriptive-stats (g/random-int 1 13)) => #(and (= (:max %) 12)
+                                                   (= (:min %) 1)
+                                                   (integer? (:max %))
+                                                   (integer? (:min %)))
+  (descriptive-stats (g/random-int -5 -2)) => #(and (= (:max %) -3)
+                                                    (= (:min %) -5)
+                                                    (integer? (:max %))
+                                                    (integer? (:min %))))
+
+(fact "On random-uniform" :slow
+  (descriptive-stats (g/random-uniform)) => #(and (< 0.95 (:max %) 1.00)
+                                                  (< 0.00 (:min %) 0.05)
+                                                  (double? (:max %))
+                                                  (double? (:min %)))
+  (descriptive-stats (g/random-uniform -0.5 -1.0)) => #(and (< -0.55 (:max %) -0.50)
+                                                            (< -1.00 (:min %) -0.95)
+                                                            (double? (:max %))
+                                                            (double? (:min %))))
+
 (fact "On random-choice" :slow
   (-> (g/table->dataset (mapv vector (range 100)) [:idx])
       (g/with-column :rand-choice (g/random-choice [(g/lit "abc")