<commit_before>(ns onyx.peer.transform
  (:require [clojure.core.async :refer [chan go alts!! close! >!] :as async]
            [onyx.peer.pipeline-extensions :as p-ext]
            [onyx.extensions :as extensions]
            [taoensso.timbre :refer [info]]
            [dire.core :refer [with-post-hook!]]))

(defn read-batch [queue consumers batch-size timeout]
  (let [consumer-chs (take (count consumers) (repeatedly #(chan 1)))]
    (doseq [[c consumer-ch] (map vector consumers consumer-chs)]
      (go (loop []
            (when-let [m (.receive c)]
              (extensions/ack-message queue m)
              (>! consumer-ch m)
              (recur)))))
    (let [chs (conj consumer-chs (async/timeout timeout))
          rets (doall (repeatedly batch-size #(first (alts!! chs))))]
      (doseq [ch chs] (close! ch))
      (filter identity rets))))

(defn decompress-segment [queue message]
  (let [segment (extensions/read-message queue message)]
    (read-string segment)))

(defn apply-fn [task segment]
  (let [user-ns (symbol (name (namespace (:onyx/fn task))))
        user-fn (symbol (name (:onyx/fn task)))]
    ((ns-resolve user-ns user-fn) segment)))

(defn compress-segment [segment]
  (pr-str segment))

(defn write-batch [queue session producers msgs]
  (for [p producers msg msgs]
    (extensions/produce-message queue p session msg)))

(defn read-batch-shim
  [{:keys [queue session ingress-queues batch-size timeout]}]
  (let [consumers (map (partial extensions/create-consumer queue session) ingress-queues)
        batch (read-batch queue consumers batch-size timeout)]
    {:batch batch :consumers consumers}))

(defn decompress-batch-shim [{:keys [queue batch]}]
  (let [decompressed-msgs (map (partial decompress-segment queue) batch)]
    {:decompressed decompressed-msgs}))

(defn apply-fn-shim [{:keys [decompressed task catalog]}]
  (let [task (first (filter (fn [entry] (= (:onyx/name entry) task)) catalog))
        results (map (partial apply-fn task) decompressed)]
    {:results results}))

(defn compress-batch-shim [{:keys [results]}]
  (let [compressed-msgs (map compress-segment results)]
    {:compressed compressed-msgs}))

(defn write-batch-shim [{:keys [queue egress-queues session compressed]}]
  (let [producers (map (partial extensions/create-producer queue session) egress-queues)
        batch (write-batch queue session producers compressed)]
    {:producers producers}))

(defmethod p-ext/read-batch :default
  [event] (read-batch-shim event))
<commit_msg>Force evaluation of for. Fixes phase 1.<commit_after>(ns onyx.peer.transform
  (:require [clojure.core.async :refer [chan go alts!! close! >!] :as async]
            [onyx.peer.pipeline-extensions :as p-ext]
            [onyx.extensions :as extensions]
            [taoensso.timbre :refer [info]]
            [dire.core :refer [with-post-hook!]]))

(defn read-batch [queue consumers batch-size timeout]
  (let [consumer-chs (take (count consumers) (repeatedly #(chan 1)))]
    (doseq [[c consumer-ch] (map vector consumers consumer-chs)]
      (go (loop []
            (when-let [m (.receive c)]
              (extensions/ack-message queue m)
              (>! consumer-ch m)
              (recur)))))
    (let [chs (conj consumer-chs (async/timeout timeout))
          rets (doall (repeatedly batch-size #(first (alts!! chs))))]
      (doseq [ch chs] (close! ch))
      (filter identity rets))))

(defn decompress-segment [queue message]
  (let [segment (extensions/read-message queue message)]
    (read-string segment)))

(defn apply-fn [task segment]
  (let [user-ns (symbol (name (namespace (:onyx/fn task))))
        user-fn (symbol (name (:onyx/fn task)))]
    ((ns-resolve user-ns user-fn) segment)))

(defn compress-segment [segment]
  (pr-str segment))

(defn write-batch [queue session producers msgs]
  (dorun
   (for [p producers msg msgs]
     (extensions/produce-message queue p session msg))))

(defn read-batch-shim
  [{:keys [queue session ingress-queues batch-size timeout]}]
  (let [consumers (map (partial extensions/create-consumer queue session) ingress-queues)
        batch (read-batch queue consumers batch-size timeout)]
    {:batch batch :consumers consumers}))

(defn decompress-batch-shim [{:keys [queue batch]}]
  (let [decompressed-msgs (map (partial decompress-segment queue) batch)]
    {:decompressed decompressed-msgs}))

(defn apply-fn-shim [{:keys [decompressed task catalog]}]
  (let [task (first (filter (fn [entry] (= (:onyx/name entry) task)) catalog))
        results (map (partial apply-fn task) decompressed)]
    {:results results}))

(defn compress-batch-shim [{:keys [results]}]
  (let [compressed-msgs (map compress-segment results)]
    {:compressed compressed-msgs}))

(defn write-batch-shim [{:keys [queue egress-queues session compressed]}]
  (let [producers (map (partial extensions/create-producer queue session) egress-queues)
        batch (write-batch queue session producers compressed)]
    {:producers producers}))

(defmethod p-ext/read-batch :default
  [event] (read-batch-shim event))

{'fraction_non_alphanumeric': 0.1458498023715415, 'fraction_numerical': 0.0005928853754940712, 'mean_word_length': 3.5106951871657754, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '331536', 'n_tokens_mistral': 1630, 'n_tokens_neox': 1599, 'n_words': 473}