dag added (#677)

Summary:
Pull Request resolved: #677

## What

* Create the dag struct with `next_column` implemented
* add Clone + Copy to ColumnMetadata

## Why

* So that we can iterate through the columns in order

## Future improvements

* Implement Iter trait so that we can iterate on the dag
* Support iterating / getting the next column level so that you can parallelize computation of some columns

Reviewed By: gorel

Differential Revision: D34772282

fbshipit-source-id: d50a05bd21c57288bdddec218bad4ffa625a50c1

Author: jrodal98

fbpcs/kodiak/Cargo.toml

@@ -7,3 +7,4 @@ edition = "2021"
 [dependencies]
 derive_more = "0.99.3"
 log = { version = "0.4.14", features = ["kv_unstable", "kv_unstable_std"] }
+petgraph = "0.6"

fbpcs/kodiak/src/column_metadata.rs

@@ -8,7 +8,7 @@
 use crate::mpc_metric_dtype::MPCMetricDType;
 use crate::row::Row;
 
-pub trait ColumnMetadata: std::cmp::Eq + std::hash::Hash + Sized {
+pub trait ColumnMetadata: std::cmp::Eq + std::hash::Hash + Sized + Clone + Copy {
     /// Used to look up a human-readable name for this metric.
     /// Should be known at compile time, so &'static is fine.
     fn name(&self) -> &'static str;
@@ -34,7 +34,7 @@ macro_rules! column_metadata {
         $($variant:ident -> [$($deps:ident),*]),*,
     }) => {
 
-        #[derive(Debug, PartialEq, Eq, std::hash::Hash)]
+        #[derive(Copy, Clone, Debug, PartialEq, Eq, std::hash::Hash)]
         pub enum $name {
             $($variant),*
         }

fbpcs/kodiak/src/dag.rs

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use crate::column_metadata::ColumnMetadata;
+use crate::mpc_view::MPCView;
+use petgraph::graph::DiGraph;
+
+pub struct Dag<T: ColumnMetadata> {
+    graph: DiGraph<T, usize>,
+    sorted_columns: Vec<T>,
+    current_index: usize,
+}
+
+impl<T: ColumnMetadata> Dag<T> {
+    pub fn next_column(&mut self) -> Option<&T> {
+        let node = self.sorted_columns.get(self.current_index);
+        self.current_index += 1;
+        node
+    }
+
+    pub fn next_columns(&mut self) -> Option<Vec<&T>> {
+        unimplemented!("Support for per-level operation not implemented")
+    }
+
+    pub fn reset(&mut self) {
+        self.current_index = 0;
+    }
+
+    pub fn from_mpc_view(mpc_view: MPCView<T>) -> Self {
+        let graph = Self::build_graph(
+            &mpc_view.input_columns,
+            &mpc_view.helper_columns,
+            &mpc_view.metrics,
+        );
+
+        let sorted_nodes = Self::toposort_nodes(&graph);
+
+        Self {
+            graph,
+            sorted_columns: sorted_nodes,
+            current_index: 0,
+        }
+    }
+
+    fn build_graph(input_columns: &[T], helper_columns: &[T], metrics: &[T]) -> DiGraph<T, usize> {
+        let mut graph = DiGraph::new();
+
+        // 1. add the metric column data to the graph
+        // 2. insert a mapping from node -> graph index
+        // 3. repeat for every column
+        let node_to_index = input_columns
+            .iter()
+            .chain(helper_columns.iter().chain(metrics.iter()))
+            .fold(std::collections::HashMap::new(), |mut acc, &node| {
+                let i = graph.add_node(node.clone());
+                acc.insert(node, i);
+                acc
+            });
+
+        for to_node in input_columns
+            .iter()
+            .chain(helper_columns.iter().chain(metrics.iter()))
+        {
+            let to_index = node_to_index
+                .get(to_node)
+                .unwrap_or_else(|| panic!("Column {} was not found in the graph", to_node.name()));
+            for from_node in to_node.dependencies().iter() {
+                let from_index = node_to_index.get(from_node).unwrap_or_else(|| {
+                    panic!("Column {} was not found in the graph", from_node.name())
+                });
+                // 0 is the weight
+                graph.add_edge(*from_index, *to_index, 0);
+            }
+        }
+
+        graph
+    }
+
+    fn toposort_nodes(graph: &DiGraph<T, usize>) -> Vec<T> {
+        match petgraph::algo::toposort(&graph, None) {
+            Ok(order) => order
+                .into_iter()
+                .map(|node_index| *graph.node_weight(node_index).unwrap())
+                .collect(),
+            Err(_e) => panic!("Cycle detected in graph"),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::column_metadata;
+    use crate::column_metadata::ColumnMetadata;
+    use crate::dag::Dag;
+    use crate::mpc_metric_dtype::MPCMetricDType;
+    use crate::mpc_view::MPCView;
+    use crate::row::Row;
+
+    column_metadata! {
+        TestEnum {
+            Variant1 -> [],
+            Variant2 -> [Variant1],
+            Variant3 -> [Variant1],
+            Variant4 -> [Variant2, Variant3],
+        }
+    }
+
+    impl TestEnum {
+        fn from_row(&self, _r: &Row<Self>) -> MPCMetricDType {
+            panic!("Undefined for test");
+        }
+        fn aggregate<I: Iterator<Item = Row<Self>>>(&self, _rows: I) -> MPCMetricDType {
+            panic!("Undefined for test");
+        }
+    }
+
+    fn get_mpc_view() -> MPCView<TestEnum> {
+        MPCView::new(
+            vec![TestEnum::Variant1],
+            vec![TestEnum::Variant2, TestEnum::Variant3],
+            vec![TestEnum::Variant4],
+            vec![],
+        )
+    }
+
+    #[test]
+    fn dag_next_node() {
+        let mpc_view = get_mpc_view();
+        let mut dag = Dag::from_mpc_view(mpc_view);
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant1));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant3));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant2));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant4));
+        assert_eq!(dag.next_column(), None);
+        dag.reset();
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant1));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant3));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant2));
+        assert_eq!(dag.next_column(), Some(&TestEnum::Variant4));
+    }
+}

fbpcs/kodiak/src/lib.rs

@@ -6,6 +6,7 @@
  */
 
 pub mod column_metadata;
+pub mod dag;
 pub mod execution_config;
 pub mod input_reader;
 pub mod metric_config;

fbpcs/kodiak/src/mpc_view.rs

@@ -8,9 +8,9 @@
 use crate::column_metadata::ColumnMetadata;
 
 pub struct MPCView<T: ColumnMetadata> {
-    input_columns: Vec<T>,
-    helper_columns: Vec<T>,
-    metrics: Vec<T>,
+    pub input_columns: Vec<T>,
+    pub helper_columns: Vec<T>,
+    pub metrics: Vec<T>,
     grouping_sets: Vec<Vec<T>>,
 }