Lazy partition-by - clojure

I have a source of items and want to separately process runs of them having the same value of a key function. In Python this would look like
for key_val, part in itertools.groupby(src, key_fn):
process(key_val, part)
This solution is completely lazy, i.e. if process doesn't try to store contents of entire part, the code would run in O(1) memory.
Clojure solution
(doseq [part (partition-by key-fn src)]
(process part))
is less lazy: it realizes each part completely. The problem is, src might have very long runs of items with the same key-fn value and realizing them might lead to OOM.
I've found this discussion where it's claimed that the following function (slightly modified for naming consistency inside post) is lazy enough
(defn lazy-partition-by [key-fn coll]
(lazy-seq
(when-let [s (seq coll)]
(let [fst (first s)
fv (key-fn fst)
part (lazy-seq (cons fst (take-while #(= fv (key-fn %)) (next s))))]
(cons part (lazy-partition-by key-fn (drop-while #(= fv (key-fn %)) s)))))))
However, I don't understand why it doesn't suffer from OOM: both parts of the cons cell hold a reference to s, so while process consumes part, s is being realized but not garbage collected. It would become eligible for GC only when drop-while traverses part.
So, my questions are:
Am I correct about lazy-partition-by not being lazy enough?
Is there an implementation of partition-by with guaranteed memory requirements, provided I don't hold any references to the previous part by the time I start realizing the next one?
EDIT:
Here's a lazy enough implementation in Haskell:
lazyPartitionBy :: Eq b => (a -> b) -> [a] -> [[a]]
lazyPartitionBy _ [] = []
lazyPartitionBy keyFn xl#(x:_) = let
fv = keyFn x
(part, rest) = span ((== fv) . keyFn) xl
in part : lazyPartitionBy keyFn rest
As can be seen from span implementation, part and rest implicitly share state. I wonder if this method could be translated into Clojure.

The rule of thumb that I use in these scenarios (ie, those in which you want a single input sequence to produce multiple output sequences) is that, of the following three desirable properties, you can generally have only two:
Efficiency (traverse the input sequence only once, thus not hold its head)
Laziness (produce elements only on demand)
No shared mutable state
The version in clojure.core chooses (1,3), but gives up on (2) by producing an entire partition all at once. Python and Haskell both choose (1,2), although it's not immediately obvious: doesn't Haskell have no mutable state at all? Well, its lazy evaluation of everything (not just sequences) means that all expressions are thunks, which start out as blank slates and only get written to when their value is needed; the implementation of span, as you say, shares the same thunk of span p xs' in both of its output sequences, so that whichever one needs it first "sends" it to the result of the other sequence, effecting the action at a distance that's necessary to preserve the other nice properties.
The alternative Clojure implementation you linked to chooses (2,3), as you noted.
The problem is that for partition-by, declining either (1) or (2) means that you're holding the head of some sequence: either the input or one of the outputs. So if you want a solution where it's possible to handle arbitrarily large partitions of an arbitrarily large input, you need to choose (1,2). There are a few ways you could do this in Clojure:
Take the Python approach: return something more like an iterator than a seq - seqs make stronger guarantees about non-mutation, and promise that you can safely traverse them multiple times, etc etc. If instead of a seq of seqs you return an iterator of iterators, then consuming items from any one iterator can freely mutate or invalidate the other(s). This guarantees consumption happens in order and that memory can be freed up.
Take the Haskell approach: manually thunk everything, with lots of calls to delay, and require the client to call force as often as necessary to get data out. This will be a lot uglier in Clojure, and will greatly increase your stack depth (using this on a non-trivial input will probably blow the stack), but it is theoretically possible.
Write something more Clojure-flavored (but still quite unusual) by having a few mutable data objects that are coordinated between the output seqs, each updated as needed when something is requested from any of them.
I'm pretty sure any of these three approaches are possible, but to be honest they're all pretty hard and not at all natural. Clojure's sequence abstraction just doesn't make it easy to produce a data structure that's what you'd like. My advice is that if you need something like this and the partitions may be too large to fit comfortably, you just accept a slightly different format and do a little more bookkeeping yourself: dodge the (1,2,3) dilemma by not producing multiple output sequences at all!
So instead of ((2 4 6 8) (1 3 5) (10 12) (7)) being your output format for something like (partition-by even? [2 4 6 8 1 3 5 10 12 7]), you could accept a slightly uglier format: ([::key true] 2 4 6 8 [::key false] 1 3 5 [::key true] 10 12 [::key false] 7). This is neither hard to produce nor hard to consume, although it is a bit lengthy and tedious to write out.
Here is one reasonable implementation of the producing function:
(defn lazy-partition-by [f coll]
(lazy-seq
(when (seq coll)
(let [x (first coll)
k (f x)]
(list* [::key k] x
((fn part [k xs]
(lazy-seq
(when (seq xs)
(let [x (first xs)
k' (f x)]
(if (= k k')
(cons x (part k (rest xs)))
(list* [::key k'] x (part k' (rest xs))))))))
k (rest coll)))))))
And here's how to consume it, first defining a generic reduce-grouped that hides the details of the grouping format, and then an example function count-partition-sizes to output the key and size of each partition without keeping any sequences in memory:
(defn reduce-grouped [f init groups]
(loop [k nil, acc init, coll groups]
(if (empty? coll)
acc
(if (and (coll? (first coll)) (= ::key (ffirst coll)))
(recur (second (first coll)) acc (rest coll))
(recur k (f k acc (first coll)) (rest coll))))))
(defn count-partition-sizes [f coll]
(reduce-grouped (fn [k acc _]
(if (and (seq acc) (= k (first (peek acc))))
(conj (pop acc) (update-in (peek acc) [1] inc))
(conj acc [k 1])))
[] (lazy-partition-by f coll)))
user> (lazy-partition-by even? [2 4 6 8 1 3 5 10 12 7])
([:user/key true] 2 4 6 8 [:user/key false] 1 3 5 [:user/key true] 10 12 [:user/key false] 7)
user> (count-partition-sizes even? [2 4 6 8 1 3 5 10 12 7])
[[true 4] [false 3] [true 2] [false 1]]
Edit: Looking at it again, I'm not really convinced my reduce-grouped is much more useful than (reduce f init (map g xs)), since it doesn't really give you any clear indication of when the key changes. So if you do need to know when a group changes, you'll want a smarter abstraction, or to use my original lazy-partition-by with nothing "clever" wrapping it.

Although this question evokes very interesting contemplation about language design, the practical problem is you want to process on partitions in constant memory. And the practical problem is resolvable with a little inversion.
Rather than processing over the result of a function that returns a sequence of partitions, pass the processing function into the function that produces the partitions. Then, you can control state in a contained manner.
First we'll provide a way to fuse together the consumption of the sequence with the state of the tail.
(defn fuse [coll wick]
(lazy-seq
(when-let [s (seq coll)]
(swap! wick rest)
(cons (first s) (fuse (rest s) wick)))))
Then a modified version of partition-by
(defn process-partition-by [processfn keyfn coll]
(lazy-seq
(when (seq coll)
(let [tail (atom (cons nil coll))
s (fuse coll tail)
fst (first s)
fv (keyfn fst)
pred #(= fv (keyfn %))
part (take-while pred s)
more (lazy-seq (drop-while pred #tail))]
(cons (processfn part)
(process-partition-by processfn keyfn more))))))
Note: For O(1) memory consumption processfn must be an eager consumer! So while (process-partition-by identity key-fn coll) is the same as (partition-by key-fn coll), because identity does not consume the partition, the memory consumption is not constant.
Tests
(defn heavy-seq []
;adjust payload for your JVM so only a few fit in memory
(let [payload (fn [] (long-array 20000000))]
(map #(vector % (payload)) (iterate inc 0))))
(defn my-process [s] (reduce + (map first s)))
(defn test1 []
(doseq [part (partition-by #(quot (first %) 10) (take 50 (heavy-seq)))]
(my-process part)))
(defn test2 []
(process-partition-by
my-process #(quot (first %) 20) (take 200 (heavy-seq))))
so.core=> (test1)
OutOfMemoryError Java heap space [trace missing]
so.core=> (test2)
(190 590 990 1390 1790 2190 2590 2990 3390 3790)

Am I correct about lazy-partition-by not being lazy enough?
Well, there's a difference between laziness and memory usage. A sequence can be lazy and still require lots of memory - see for instance the implementation of clojure.core/distinct, which uses a set to remember all the previously observed values in the sequence. But yes, your analysis of the memory requirements of lazy-partition-by is correct - the function call to compute the head of the second partition will retain the head of the first partition, which means that realizing the first partition causes it to be retained in-memory. This can be verified with the following code:
user> (doseq [part (lazy-partition-by :a
(repeatedly
(fn [] {:a 1 :b (long-array 10000000)})))]
(dorun part))
; => OutOfMemoryError Java heap space
Since neither doseq nor dorun retains the head, this would simply run forever if lazy-partition-by were O(1) in memory.
Is there an implementation of partition-by with guaranteed memory requirements, provided I don't hold any references to the previous part by the time I start realizing the next one?
It would be very difficult, if not impossible, to write such an implementation in a purely functional manner that would work for the general case. Consider that a general lazy-partition-by implementation cannot make any assumptions about when (or if) a partition is realized. The only guaranteed correct way of finding the start of the second partition, short of introducing some nasty bit of statefulness to keep track of how much of the first partition has been realized, is to remember where the first partition began and scan forward when requested.
For the special case where you're processing records one at a time for side effects and want them grouped by key (as is implied by your use of doseq above), you might consider something along the lines of a loop/recur which maintains a state and re-sets it when the key changes.

Related

What scope should calls to lazy-seq have?

I'm writing a lazy implementation of the Recamán's Sequence, and ran into some confusion regarding where calls to lazy-seq should happen.
This first version I came up with this morning was:
(defn lazy-recamans-sequence []
(let [f (fn rec [n seen last-s]
(let [back (- last-s n)
new-s (if (and (pos? back) (not (seen back)))
back
(+ last-s n))]
(lazy-seq ; Here
(cons new-s (rec (inc n) (conj seen new-s) new-s)))))]
(f 0 #{} 0)))
Then I realized that my placement of lazy-seq was kind of arbitrary, and that it could be placed higher to wrap more of the computations:
(defn lazy-recamans-sequence2 []
(let [f (fn rec [n seen last-s]
(lazy-seq ; Here
(let [back (- last-s n)
new-s (if (and (pos? back) (not (seen back)))
back
(+ last-s n))]
(cons new-s (rec (inc n) (conj seen new-s) new-s)))))]
(f 0 #{} 0)))
Then I looked back on a review that someone gave me last night:
(defn recaman []
(letfn [(tail [previous n seen]
(let [nx (if (and (> previous n) (not (seen (- previous n))))
(- previous n)
(+ previous n))]
; Here, inside "cons"
(cons nx (lazy-seq (tail nx (inc n) (conj seen nx))))))]
(tail 0 0 #{})))
And they have theirs inside of the call to cons!
Thinking this over, it seems like it wouldn't make a difference. With a broader scope (like the second version), more code is inside the explicit function that's passed to LazySeq. With a narrower scope however, the function itself may be smaller, but since the passed function involves a recursive call, it will be executing the same code anyways.
They seem to preform nearly identically and give the same answers. Is there any reason to prefer placing lazy-seq in one place over another? Is this simply a stylistic choice, or can this have actual repercussions?
In the first two examples the lazy-seq wraps the cons call. This means that when you generate call the function you return a lazy sequence immediately without calculating the first item of the sequence.
In the first example the let expression is still outside of lazy-seq so the value of the first item is calculated immediately but the returned sequence is still lazy and not realized.
The second example is similar to the first. The lazy-seq wraps the cons cell and also the let block. This means that the function will return immediatetly and the value of the first item is calculated only when the caller starts to consume the lazy sequence.
In the third example the value of the first item in the list is calculated immediately and only the tail of the returned sequence is lazy.
Is there any reason to prefer placing lazy-seq in one place over another?
It depends on what you want to achieve. Do you want to return a sequence immediately without calculating any values? In this case make the scope of lazy-seq as broad as possible. Otherwise try to restrict the scope of lazy-seq to calculate only the tail part of the sequence.
When I was first learning Clojure, I was a bit confused by the many possible choices of lazy-seq constructs, the lack of clarity in terms of which construct to choose, and the somewhat vague explanation for how lazy-seq creates laziness in the first place (it is implemented as a Java class of ~240 lines).
To reduce repetition and keep things as simple as possible, I created the lazy-cons macro. It is used like so:
(defn lazy-countdown [n]
(when (<= 0 n)
(lazy-cons n (lazy-countdown (dec n)))))
(deftest t-all
(is= (lazy-countdown 5) [5 4 3 2 1 0] )
(is= (lazy-countdown 1) [1 0] )
(is= (lazy-countdown 0) [0] )
(is= (lazy-countdown -1) nil ))
This version does realize the initial value n immediately.
I never worry about chunking (typically batches of 32) or trying to precisely control the number of elements realized in a lazy sequence. IMHO, if you need fine-grained control such as this, it is better to use an explicit loop than to make assumptions on the timing of realizations in a lazy sequence.

How do Clojure transducers work under the hood?

Reading all over (except the clojure source code) it is somewhat hard to fathom how do transducers avoid the usage of intermediate collections, which is supposed to make them more lean and performant.
A related question arises as to whether or not, they assume, that each input transformation is applied to each element of its input independently from the other elements of it, a limitation that may exist if transducers were to work by squashing the input transformations on the input collection ― element-by-element.
Do they inspect the code of their input functions to determine how to interweave them such that they yield the correct result of their composition?
Can you please detail how do transducers in clojure work under the hood, in those regards?
A related question arises as to whether or not, they assume, that each
input transformation is applied to each element of its input
independently from the other elements of it
They are named transducers because they may have (implicit) state.
A transducer is a function that takes a reducing function and returns a reducing function.
A reducing function is a function that expects two parameters: an accumulator and an item and returns an updated accumulator.
This is the reducing function which holds the mutable state (if any).
To get transducers you have to understand they work in two times: composition-time then computation-time. That's why they are functions returning functions.
Let's start with an easy reducing function: conj.
The transducer returned by (map inc) is (fn [rf] (fn [acc x] (rf acc (inc x)))). When called with conj it returns a function equivalent to (fn [acc x] (conj acc (inc x))).
The transducer returned by (filter odd?) is (fn [rf] (fn [acc x] (if (odd? x) (rf acc x) acc))). When called with conj it returns a function equivalent to (fn [acc x] (if (odd? x) (conj acc x) acc))). This one is interesting because rf (the downstream reducing function is sometimes short-circuited).
If you want to chain these two transducers you just do (comp (map inc) (filter odd?)) if you pass conj to this composite transducer, (filter odd?) is going to be the first to wrap conj (because comp applies functions from right to left). Then the resulting filtered-rffunction is passed to (map inc) which yields a function equiavlent to:
(fn [acc x] (filtered-rf acc (inc x))) where filtered-rf is (fn [acc x] (if (odd? x) (conj acc x) acc))). If you inline filtered-rf you get: (fn [acc x] (let [x+1 (inc x)] (if (odd? x+1) (conj acc x+1) acc))).
As you may see no intermediate collection or sequence is allocated.
For stateful transducers that's the same story except that reducing functions have mutable state (as little as possible and avoid keeping all previous items in it): usually a volatile box (see volatile!) or a mutable Java object.
You may also have remarked that in the example items are first mapped then filtered: computations are applied from left to right which seems in contradiction to comp. This is not: remember comp here composes transducers, fns that wraps reducing fns. So at compoition time the wrapping occurs right to left (conj wrapped by the "filtering rf" then by the "mapping rf") but at computation time the wrapping layers are traversed inwards: map, filter and then conj.
There are finicky implementation details to know to implement your own transducers (reduced, init and completion arities) but the general idea is the one exposed above.

lazy-seq -- cons outside or in

Should cons be inside (lazy-seq ...)
(def lseq-in (lazy-seq (cons 1 (more-one))))
or out?
(def lseq-out (cons 1 (lazy-seq (more-one))))
I noticed
(realized? lseq-in)
;;; ⇒ false
(realized? lseq-out)
;;; ⇒ <err>
;;; ClassCastException clojure.lang.Cons cannot be cast to clojure.lang.IPending clojure.core/realized? (core.clj:6773)
All the examples on the clojuredocs.org use "out".
What are the tradeoffs involved?
You definitely want (lazy-seq (cons ...)) as your default, deviating only if you have a clear reason for it. clojuredocs.org is fine, but the examples are all community-provided and I would not call them "the docs". Of course, a consequence of how it's built is that the examples tend to get written by people who just learned how to use the construct in question and want to help out, so many of them are poor. I would refer instead to the code in clojure.core, or other known-good code.
Why should this be the default? Consider these two implementations of map:
(defn map1 [f coll]
(when-let [s (seq coll)]
(cons (f (first s))
(lazy-seq (map1 f (rest coll))))))
(defn map2 [f coll]
(lazy-seq
(when-let [s (seq coll)]
(cons (f (first s))
(map2 f (rest coll))))))
If you call (map1 prn xs), then an element of xs will be realized and printed immediately, even if you never intentionally realize an element of the resulting mapped sequence. map2, on the other hand, immediately returns a lazy sequence, delaying all its work until an element is requested.
With cons inside lazy-seq, the evaluation of the expression for the first element of your seq gets deferred; with cons on the outside, it's done right away and only the construction of the "rest" part of the seq is deferred. (So (rest lseq-out) will be a lazy seq.)
Thus, if computing the first element is expensive and it might not be needed at all, putting cons inside lazy-seq makes more sense. If the initial element is supplied to the lazy seq producer as an argument, it may make more sense to use cons on the outside (this is the case with clojure.core/iterate). Otherwise it doesn't make that much of a difference. (The overhead of creating a lazy seq object at the start is negligible.)
Clojure itself uses both approaches (although in the majority of cases lazy-seq wraps the whole seq-producing expression, which may not necessarily start with cons).

Is there a way to control the number of threads used with pmap?

I am just doing some performance testing with clojure using pmap and I would like to be able to control the number of threads being used with pmap. I know when using something like OpenMP one can set the number of threads using omp_set_num_threads(). I was wondering if there would be anything similar in clojure.
Here's code for pmap:
(defn pmap
"Like map, except f is applied in parallel. Semi-lazy in that the
parallel computation stays ahead of the consumption, but doesn't
realize the entire result unless required. Only useful for
computationally intensive functions where the time of f dominates
the coordination overhead."
([f coll]
(let [n (+ 2 (.. Runtime getRuntime availableProcessors))
rets (map #(future (f %)) coll)
step (fn step [[x & xs :as vs] fs]
(lazy-seq
(if-let [s (seq fs)]
(cons (deref x) (step xs (rest s)))
(map deref vs))))]
(step rets (drop n rets))))
As you can see, pmap takes all available processors and uses them cyclically. So, no, there's no way to set the number of threads... but you can always write your own pmap, which will provide such functionality.

How to break out from nested doseqs

I have a question regarding nested doseq loops. In the start function, once I find an answer I set the atom to true, so that the outer loop validation with :while fails. However it seems that it doesn't break it, and the loops keep on going. What's wrong with it?
I am also quite confused with the usage of atoms, refs, agents (Why do they have different names for the update functions when then the mechanism is almost the same?) etc.
Is it okay to use an atom in this situation as a flag? Obviously I need a a variable like object to store a state.
(def pentagonal-list (map (fn [a] (/ (* a (dec (* 3 a))) 2)) (iterate inc 1)))
(def found (atom false))
(defn pentagonal? [a]
(let [y (/ (inc (Math/sqrt (inc (* 24 a)))) 6)
x (mod (* 10 y) 10)]
(if (zero? x)
true
false)))
(defn both-pent? [a b]
(let [sum (+ b a)
diff (- a b)]
(if (and (pentagonal? sum) (pentagonal? diff))
true
false)))
(defn start []
(doseq [x pentagonal-list :while (false? #found)]
(doseq [y pentagonal-list :while (<= y x)]
(if (both-pent? x y)
(do
(reset! found true)
(println (- x y)))))))
Even once the atom is set to true, your version can't stop running until the inner doseq finishes (until y > x). It will terminate the outer loop once the inner loop finishes though. It does terminate eventually when I run it. Not sure what you're seeing.
You don't need two doseqs to do this. One doseq can handle two seqs at once.
user> (doseq [x (range 0 2) y (range 3 6)] (prn [x y]))
[0 3]
[0 4]
[0 5]
[1 3]
[1 4]
[1 5]
(The same is true of for.) There is no mechanism for "breaking out" of nested doseqs that I know of, except throw/catch, but that's rather un-idiomatic. You don't need atoms or doseq for this at all though.
(def answers (filter (fn [[x y]] (both-pent? x y))
(for [x pentagonal-list
y pentagonal-list :while (<= y x)]
[x y])))
Your code is very imperative in style. "Loop over these lists, then test the values, then print something, then stop looping." Using atoms for control like this is not very idiomatic in Clojure.
A more functional way is to take a seq (pentagonal-list) and wrap it in functions that turn it into other seqs until you get a seq that gives you what you want. First I use for to turn two copies of this seqs into one seq of pairs where y <= x. Then I use filter to turn that seq into one that filters out the values we don't care about.
filter and for are lazy, so this will stop running once it finds the first valid value, if one is all you want. This returns the two numbers you want, and then you can subtract them.
(apply - (first answers))
Or you can further wrap the function in another map to calculate the differences for you.
(def answers2 (map #(apply - %) answers))
(first answers2)
There are some advantages to programming functionally in this way. The seq is cached (if you hold onto the head like I do here), so once a value is calculated, it remembers it and you can access it instantly from then on. Your version can't be run again without resetting the atom, and then would have to re-calculate everything. With my version you can (take 5 answers) to get the first 5 results, or map over the result to do other things if you want. You can doseq over it and print the values. Etc. etc.
I'm sure there are other (maybe better) ways to do this still without using an atom. You should usually avoid mutating references unless it's 100% necessary in Clojure.
The function names for altering atoms/agents/refs are different probably because the mechanics are not the same. Refs are synchronous and coordinated via transactions. Agents are asynchronous. Atoms are synchronous and uncoordinated. They all kind of "change a reference", and probably some kind of super-function or macro could wrap them all in one name, but that would obscure the fact that they're doing drastically different things under the hood. Explaining the differences fully is probably beyond the scope of an SO post to explain, but http://clojure.org fully explains all of the nuances of the differences.