Lazy partition-by

Question

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:

1. Am I correct about lazy-partition-by not being lazy enough?
2. 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.

Answer 1

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:

1. Efficiency (traverse the input sequence only once, thus not hold its head)
2. Laziness (produce elements only on demand)
3. 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:

1. 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.
2. 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.
3. 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.

Answer 2

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)

Answer 3

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.