   (poll! [this _]
     (if (and iter (.hasNext ^java.util.Iterator iter))
       (let [rec ^ConsumerRecord (.next ^java.util.Iterator iter)
-            new-offset (if rec (.offset rec) offset)
             deserialized (some-> rec segment-fn)]
-        (if (= :done deserialized)
-          (do (set! drained true)
-              nil)
-          (do
-           (.set ^AtomicLong read-offset new-offset)
-           (set! offset new-offset) 
-           deserialized))) 
+        (cond (= :done deserialized)
+              (do (set! drained true)
+                  nil)
+              deserialized
+              (let [new-offset (.offset rec)
+                    part (.partition rec)]
+                (.set ^AtomicLong read-offset new-offset)
+                (set! partition->offset (assoc partition->offset part new-offset))
+                deserialized)))
       (do (set! iter 
                 (.iterator ^ConsumerRecords 
                            (.poll ^Consumer (.consumer ^FranzConsumer consumer) 

  (let [policy (:kafka/offset-reset task-map)
        start-offsets (:kafka/start-offsets task-map)]
    (doseq [kpartition kpartitions]
      (cond checkpoint
            (do
              (info log-prefix "Seeking to checkpointed offset at:" (inc checkpoint))
              (seek-to-offset! consumer {:topic topic :partition kpartition} (inc checkpoint)))

            start-offsets
            (let [offset (get start-offsets kpartition)]


  p/Checkpointed
  (checkpoint [this]
    partition->offset)

  (recover! [this replica-version checkpoint]
    (set! drained false)
    (set! iter nil)
    (set! partition->offset checkpoint)
    (seek-offset! log-prefix consumer kpartitions task-map topic checkpoint)
    this)

