Lazy partition-by

clojure, lazy-evaluation, lazy-sequences

Solution

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)

Problem

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.

Original source