         (recur it)))))
 
 (defn cogroup-transform
-  ([specs {:keys [join-nil?] :as options}]
+  ([{:keys [join-nil?] :as options}]
    (let [opts (assoc options :label :raw-cogroup)]
      (ptransform
       :cogroup
-      [^KeyedPCollectionTuple pcolltuple]
-      (let [ordered-tags (->> pcolltuple
+      [^GroupSpecs group-specs]
+      (let [root-name (or (:name options) "cogroup")
+            pcolls (for [[idx [pcoll f {:keys [drop-nil?] :as opts}]]
+                         (map-indexed (fn [idx s] (if (instance? PCollection s)
+                                                    [idx [s nil nil]] [idx s])) (:specs group-specs))]
+                     (let [local-name (str root-name "-" (.getName pcoll))
+                           op (if f
+                                (with-keys f {:name (str local-name "-group-by")} pcoll)
+                                pcoll)]
+                       (if drop-nil?
+                         (dfilter (if f
+                                    (fn [^KV kv] (not (nil? (.getKey kv))))
+                                    (fn [v] (not (nil? v))))
+                                  {:name (str local-name "-drop-nil")
+                                   :without-coercion-to-clj true}
+                                  op)
+                         op)))
+            pcolltuple (make-keyed-pcollection-tuple pcolls)
+            ordered-tags (->> pcolltuple
                               (.getKeyedCollections)
                               (map #(.getTupleTag %))
                               (sort-by #(.getId %)))
             rel (apply-transform pcolltuple (CoGroupByKey/create) base-schema opts)
+            required-set (->> (:specs group-specs)
+                              (map-indexed (fn [idx [_ _ {:keys [type]}]]
+                                             (when (= type :required) idx)))
+                              (remove nil?)
+                              (into #{}))
+            required-count (count required-set)
             final-rel (pardo
                        (fn [^DoFn$ProcessContext c]
                          (let [^KV kv (.element c)
                                k (.getKey kv)
-                               ^CoGbkResult raw-values (.getValue kv)
-                               required-list (->> specs
-                                                  (map-indexed
-                                                   (fn [idx {:keys [type]}]
-                                                     (when (= type :required) idx)))
-                                                  (remove nil?))
-                               required-count (count required-list)]
+                               ^CoGbkResult raw-values (.getValue kv)]
                            ;; skip if a required part of the group is empty
                            (when-not (some identity
-                                           (map (fn [tag {:keys [type]}]
-                                                  (if (= type :required)
-                                                    (not (.hasNext (.iterator (.getAll raw-values tag))))
-                                                    false))
-                                                ordered-tags specs))
+                                           (map-indexed (fn [idx tag]
+                                                          (if (required-set idx)
+                                                            (not (-> (.getAll raw-values tag)
+                                                                     (.iterator)
+                                                                     (.hasNext)))
+                                                            false))
+                                                        ordered-tags))
                              (if (and (not join-nil?) (nil? k))
                                (cond
-                                 (= required-count 0) (doseq [[idx tag] (map-indexed (fn [idx tag] [idx tag]) ordered-tags)]
-                                                        (greedy-emit-cogbkresult raw-values (count ordered-tags) idx tag c))
-                                 (= required-count 1) (greedy-emit-cogbkresult raw-values
-                                                                               (count ordered-tags)
-                                                                               (first required-list)
-                                                                               (nth ordered-tags
-                                                                                    (first required-list))
-                                                                               c)
+                                 (= required-count 0) (doseq [[idx tag]
+                                                              (map-indexed
+                                                               (fn [idx tag] [idx tag]) ordered-tags)]
+                                                        (greedy-emit-cogbkresult
+                                                         raw-values (count ordered-tags) idx tag c))
+                                 (= required-count 1) (greedy-emit-cogbkresult
+                                                       raw-values
+                                                       (count ordered-tags)
+                                                       (first required-set)
+                                                       (nth ordered-tags
+                                                            (first required-set))
+                                                       c)
                                  :else nil)
                                (let [values (mapv
                                              (fn [tag]

       tr))
   ptransform schema))

(defn apply-transform
  "apply the PTransform to the given Pcoll applying options according to schema."
  [pcoll ^PTransform transform schema


Example:
```
(let [pcoll1 (ds/generate-input [[:a 11] [:b 12] [:a 13]])
      pcoll2 (ds/generate-input [[:a 21] [:c 22]])]
  (ds/cogroup [{:type :required} {:type :optional}] {:name :my-cogroup} [pcoll1 pcoll2]))
;; returns:
;; '([:a ['(11 13) '(21)]]
