compute: Emit TopK monotonic and basic output as the columnar edge

src/compute/src/render/top_k.rs

@@ -36,6 +36,7 @@ use mz_ore::cast::CastFrom;
 use mz_ore::soft_assert_or_log;
 use mz_repr::fixed_length::ExtendDatums;
 use mz_repr::{Datum, DatumVec, Diff, ReprScalarType, Row, SharedRow};
+use mz_timely_util::columnar::builder::ColumnBuilder;
 use mz_timely_util::columnation::ColumnationChunker;
 use mz_timely_util::operator::CollectionExt;
 use timely::Container;
@@ -293,10 +294,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
                         "requested no validation, but received error collection"
                     );
 
-                    CollectionBundle::from_collections(
-                        result.map(|(_key_hash, row)| row),
-                        err_collection,
-                    )
+                    CollectionBundle::from_edge(topk_result_to_columnar(result), err_collection)
                 }
                 TopKPlan::Basic(BasicTopKPlan {
                     group_key,
@@ -319,7 +317,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
                         ok_input, group_key, order_key, offset, limit, arity, buckets,
                     );
                     err_collection = err_collection.concat(errs);
-                    CollectionBundle::from_collections(oks, err_collection)
+                    CollectionBundle::from_edge(oks, err_collection)
                 }
             };
 
@@ -341,7 +339,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
         arity: usize,
         buckets: Vec<u64>,
     ) -> (
-        VecCollection<'s, T, Row, Diff>,
+        CollectionEdge<'s, T>,
         VecCollection<'s, T, DataflowErrorSer, Diff>,
     ) {
         let pairer = Pairer::new(1);
@@ -417,7 +415,7 @@ impl<'scope, T: crate::render::RenderTimestamp + crate::render::MaybeBucketByTim
             err_collection = errs;
         }
         (
-            collection.map(|(_key_hash, row)| row),
+            topk_result_to_columnar(collection),
             err_collection.expect("at least one stage validated its inputs"),
         )
     }
@@ -663,6 +661,36 @@ where
     }
 }
 
+/// Drops the hash-key pairing from a consolidated `(hash_key, row)` TopK result.
+///
+/// The hash key is a function of the row and the input is consolidated, so dropping the
+/// key is injective and the output has no within-batch duplicates for a consolidating
+/// builder to fold. Rows are pushed borrowed.
+///
+/// TODO: TopK renders its stages over `Vec` containers, so this encode sits at the very
+/// end of the plan. Pushing columnar containers down through `build_topk` and the
+/// monotonic path would remove it.
+fn topk_result_to_columnar<'s, T>(
+    collection: VecCollection<'s, T, (Row, Row), Diff>,
+) -> CollectionEdge<'s, T>
+where
+    T: crate::render::RenderTimestamp,
+{
+    let stream = collection
+        .inner
+        .unary::<ColumnBuilder<(Row, T, Diff)>, _, _, _>(Pipeline, "TopKUnkey", |_cap, _info| {
+            move |input, output| {
+                input.for_each(|time, data| {
+                    let mut session = output.session_with_builder(&time);
+                    for ((_key_hash, row), t, d) in data.drain(..) {
+                        session.give((&row, &t, &d));
+                    }
+                });
+            }
+        });
+    CollectionEdge::Columnar(stream.as_collection())
+}
+
 /// Build a stage of a topk reduction. Maintains the _retractions_ of the output instead of emitted
 /// rows. This has the benefit that we have to maintain state proportionally to size of the output
 /// instead of the size of the input.
@@ -1328,4 +1356,57 @@ mod tests {
             assert_eq!(key_datums[1], value_datums[0]);
         }
     }
+
+    #[mz_ore::test]
+    fn topk_result_to_columnar_drops_key() {
+        let key = Row::pack_slice(&[Datum::Int64(7)]);
+        let rows = vec![
+            (
+                (key.clone(), Row::pack_slice(&[Datum::Int32(1)])),
+                0u64,
+                Diff::ONE,
+            ),
+            (
+                (key.clone(), Row::pack_slice(&[Datum::Int32(2)])),
+                1u64,
+                Diff::ONE,
+            ),
+            // Retracts at a `(row, time)` with no insertion, so the `InputSession`'s
+            // pre-send consolidation does not cancel it out.
+            (
+                (key.clone(), Row::pack_slice(&[Datum::Int32(1)])),
+                2u64,
+                -Diff::ONE,
+            ),
+        ];
+        let mut expected: Vec<(Row, Timestamp, Diff)> = rows
+            .iter()
+            .map(|((_, v), t, d)| (v.clone(), Timestamp::from(*t), *d))
+            .collect();
+        expected.sort();
+
+        let (is_columnar, captured) = timely::execute_directly(move |worker| {
+            worker.dataflow::<Timestamp, _, _>(|scope| {
+                let (mut handle, collection) = scope.new_collection();
+                let edge = topk_result_to_columnar(collection);
+                let is_columnar = matches!(edge, CollectionEdge::Columnar(_));
+                let captured = edge.into_vec().inner.capture();
+                for (kv, time, diff) in rows {
+                    handle.update_at(kv, Timestamp::from(time), diff);
+                }
+                handle.advance_to(Timestamp::from(3u64));
+                handle.flush();
+                (is_columnar, captured)
+            })
+        });
+        assert!(is_columnar, "the TopK output must be a columnar edge");
+
+        let mut got: Vec<(Row, Timestamp, Diff)> = captured
+            .extract()
+            .into_iter()
+            .flat_map(|(_, data)| data)
+            .collect();
+        got.sort();
+        assert_eq!(got, expected);
+    }
 }