first draft fully merge group

Author: connortsui20

optd-mvp/src/memo/persistent/implementation.rs

@@ -90,12 +90,55 @@ impl PersistentMemo {
         // For every group along the path that we walked, set their parent id pointer to the root.
         // This allows for an amortized O(1) cost for `get_root_group`.
         for group in path {
-            self.update_group_parent(GroupId(group.id), root_id).await?;
+            let mut active_group = group.into_active_model();
+
+            // Update the group to point to the new parent.
+            active_group.parent_id = Set(Some(root_id.0));
+            active_group.update(&self.db).await?;
         }
 
         Ok(root_id)
     }
 
+    /// Retrieves every group ID of groups that share the same root group with the input group.
+    ///
+    /// If a group does not exist in the cycle, returns a [`MemoError::UnknownGroup`] error.
+    ///
+    /// The group records form a union-find data structure that also maintains a circular linked
+    /// list in every set that allows us to iterate over all elements in a set in linear time.
+    pub async fn get_group_set(&self, group_id: GroupId) -> OptimizerResult<Vec<GroupId>> {
+        // Iterate over the circular linked list until we reach ourselves again.
+        let base_group = self.get_group(group_id).await?;
+
+        // The only case when `next_id` is set to `None` is if the current group is a root, which
+        // means that this group is the only group in the set.
+        if base_group.next_id.is_none() {
+            assert!(base_group.parent_id.is_none());
+            return Ok(vec![group_id]);
+        }
+
+        // Iterate over the circular linked list until we see ourselves again, collecting nodes
+        // along the way.
+        let mut set = vec![];
+        let mut next_id = base_group
+            .next_id
+            .expect("next pointer cannot be null if it is in a cycle");
+        loop {
+            let curr_group = self.get_group(GroupId(next_id)).await?;
+
+            if curr_group.id == group_id.0 {
+                break;
+            }
+
+            set.push(GroupId(curr_group.id));
+            next_id = curr_group
+                .next_id
+                .expect("next pointer cannot be null if it is in a cycle");
+        }
+
+        Ok(set)
+    }
+
     /// Retrieves a [`physical_expression::Model`] given a [`PhysicalExpressionId`].
     ///
     /// If the physical expression does not exist, returns a
@@ -227,30 +270,6 @@ impl PersistentMemo {
         Ok(old_id)
     }
 
-    /// Updates / replaces a group's parent group. Optionally returns the previous parent.
-    ///
-    /// If either of the groups do not exist, returns a [`MemoError::UnknownGroup`] error.
-    pub async fn update_group_parent(
-        &self,
-        group_id: GroupId,
-        parent_id: GroupId,
-    ) -> OptimizerResult<Option<GroupId>> {
-        // First retrieve the group record.
-        let mut group = self.get_group(group_id).await?.into_active_model();
-
-        // Check that the parent group exists.
-        let _ = self.get_group(parent_id).await?;
-
-        // Update the group to point to the new parent.
-        let old_parent = group.parent_id;
-        group.parent_id = Set(Some(parent_id.0));
-        group.update(&self.db).await?;
-
-        // Note that the `unwrap` here is unwrapping the `ActiveValue`, not the `Option`.
-        let old_parent = old_parent.unwrap().map(GroupId);
-        Ok(old_parent)
-    }
-
     /// Adds a logical expression to an existing group via its ID.
     ///
     /// The caller is required to pass in a slice of [`GroupId`] that represent the child groups of
@@ -323,7 +342,7 @@ impl PersistentMemo {
             kind: Set(kind),
             hash: Set(hash),
         };
-        let _ = fingerprint::Entity::insert(fingerprint)
+        fingerprint::Entity::insert(fingerprint)
             .exec(&self.db)
             .await?;
 
@@ -513,10 +532,92 @@ impl PersistentMemo {
             kind: Set(kind),
             hash: Set(hash),
         };
-        let _ = fingerprint::Entity::insert(fingerprint)
+        fingerprint::Entity::insert(fingerprint)
             .exec(&self.db)
             .await?;
 
         Ok(Ok((GroupId(group_id), LogicalExpressionId(expr_id))))
     }
+
+    /// Merges two groups sets together.
+    ///
+    /// If either of the input groups do not exist, returns a [`MemoError::UnknownGroup`] error.
+    ///
+    /// TODO write docs.
+    /// TODO highly inefficient, need to understand metrics and performance testing.
+    /// TODO Optimization: add rank / size into data structure
+    pub async fn merge_groups(
+        &self,
+        left_group_id: GroupId,
+        right_group_id: GroupId,
+    ) -> OptimizerResult<GroupId> {
+        // Without a rank / size field, we have no way of determining which set is better to merge
+        // into the other. So we will arbitrarily choose to merge the left group into the right
+        // group here. If rank is added in the future, then merge the smaller set into the larger.
+
+        let left_root_id = self.get_root_group(left_group_id).await?;
+        let left_root = self.get_group(left_root_id).await?;
+        // A `None` next pointer means it should technically be pointing to itself.
+        let left_next = left_root.next_id.unwrap_or(left_root_id.0);
+        let mut active_left_root = left_root.into_active_model();
+
+        let right_root_id = self.get_root_group(right_group_id).await?;
+        let right_root = self.get_group(right_root_id).await?;
+        // A `None` next pointer means it should technically be pointing to itself.
+        let right_next = right_root.next_id.unwrap_or(right_root_id.0);
+        let mut active_right_root = right_root.into_active_model();
+
+        // Before we actually update the group records, We first need to generate new fingerprints
+        // for every single expression that has a child group in the left set.
+        // TODO make this more efficient, this code is doing double work from `get_group_set`.
+        let group_set_ids = self.get_group_set(left_group_id).await?;
+        let mut left_group_models = Vec::with_capacity(group_set_ids.len());
+        for &group_id in &group_set_ids {
+            left_group_models.push(self.get_group(group_id).await?);
+        }
+
+        // Retrieve every single expression that has a child group in the left set.
+        let left_group_expressions: Vec<Vec<logical_expression::Model>> = left_group_models
+            .load_many_to_many(
+                logical_expression::Entity,
+                logical_children::Entity,
+                &self.db,
+            )
+            .await?;
+
+        // Need to replace every single occurrence of groups in the set with the new root.
+        let rewrites: Vec<(GroupId, GroupId)> = group_set_ids
+            .iter()
+            .map(|&group_id| (group_id, right_root_id))
+            .collect();
+
+        // For each expression, generate a new fingerprint.
+        for model in left_group_expressions.into_iter().flatten() {
+            let expr_id = model.id;
+            let logical_expression: LogicalExpression = model.into();
+            let hash = logical_expression.fingerprint_with_rewrite(&rewrites);
+
+            let fingerprint = fingerprint::ActiveModel {
+                id: NotSet,
+                logical_expression_id: Set(expr_id),
+                kind: Set(logical_expression.kind()),
+                hash: Set(hash),
+            };
+            fingerprint::Entity::insert(fingerprint)
+                .exec(&self.db)
+                .await?;
+        }
+
+        // Update the left group root to point to the right group root.
+        active_left_root.parent_id = Set(Some(right_root_id.0));
+
+        // Swap the next pointers of each root to maintain the circular linked list.
+        active_left_root.next_id = Set(Some(right_next));
+        active_right_root.next_id = Set(Some(left_next));
+
+        active_left_root.update(&self.db).await?;
+        active_right_root.update(&self.db).await?;
+
+        Ok(right_root_id)
+    }
 }

optd-mvp/src/memo/persistent/tests.rs

@@ -191,22 +191,11 @@ async fn test_simple_group_link() {
 
     // The above tells the application that the expression already exists in the memo, specifically
     // under `existing_group`. Thus, we should link these two groups together.
-    // Here, we arbitrarily choose to link group 1 into group 2.
-    memo.update_group_parent(join_group_1, join_group_2)
-        .await
-        .unwrap();
+    memo.merge_groups(join_group_1, join_group_2).await.unwrap();
 
     let test_root_1 = memo.get_root_group(join_group_1).await.unwrap();
     let test_root_2 = memo.get_root_group(join_group_2).await.unwrap();
     assert_eq!(test_root_1, test_root_2);
 
-    // TODO(Connor)
-    //
-    // We now need to find all logical expressions that had group 1 (or whatever the root group of
-    // the set that group 1 belongs to is, in this case it is just group 1) as a child, and add a
-    // new fingerprint for each one that uses group 2 as a child instead.
-    //
-    // In order to do this, we need to iterate through every group in group 1's set.
-
     memo.cleanup().await;
 }