feat(core): new task graph without rules and merging

optd/src/core/error.rs

@@ -1,7 +1,30 @@
+use crate::memo::MemoizeError;
+use futures::channel::{mpsc, oneshot::Canceled};
+
 /// all optd-core errors, defined in there respective modules,
 /// but everyone uses this common Error type for simplicity.
 #[derive(Debug)]
 #[allow(dead_code)]
 pub enum Error {
     Placeholder,
+    Channel(Box<dyn std::error::Error + Send + Sync>),
+    Memo(MemoizeError),
+}
+
+impl From<mpsc::SendError> for Error {
+    fn from(value: mpsc::SendError) -> Self {
+        Error::Channel(Box::new(value))
+    }
+}
+
+impl From<Canceled> for Error {
+    fn from(value: Canceled) -> Self {
+        Error::Channel(Box::new(value))
+    }
+}
+
+impl From<MemoizeError> for Error {
+    fn from(value: MemoizeError) -> Self {
+        Error::Memo(value)
+    }
 }

optd/src/core/mod.rs

@@ -1,4 +1,4 @@
 pub mod bridge;
 pub mod cir;
 pub mod error;
-pub mod memo;
+pub mod optimizer;

optd/src/core/optimizer/client.rs

@@ -0,0 +1,117 @@
+use futures::{
+    SinkExt, StreamExt,
+    channel::{mpsc, oneshot},
+};
+
+use crate::{
+    core::cir::{LogicalPlan, PhysicalPlan},
+    core::error::Error,
+    memo::Memoize,
+};
+
+/// Unique identifier for a query instance.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct QueryInstanceId(pub i64);
+
+pub struct QueryInstance {
+    /// The logical plan associated with this query instance.
+    id: QueryInstanceId,
+    logical_plan: LogicalPlan,
+    client_tx: mpsc::Sender<ClientMessage>,
+    physical_plan_rx: mpsc::Receiver<PhysicalPlan>,
+}
+
+impl QueryInstance {
+    /// Gets the id of this query instance.
+    pub fn id(&self) -> QueryInstanceId {
+        self.id
+    }
+
+    /// Gets the logical plan for this query instance.
+    pub fn logical_plan(&self) -> &LogicalPlan {
+        &self.logical_plan
+    }
+
+    /// Receives the best physical plan for this query instance.
+    pub async fn recv_best_plan(&mut self) -> Result<PhysicalPlan, Error> {
+        let physical_plan = self
+            .physical_plan_rx
+            .next()
+            .await
+            .ok_or_else(|| Error::Placeholder)?;
+        Ok(physical_plan)
+    }
+}
+
+impl Drop for QueryInstance {
+    fn drop(&mut self) {
+        let message = ClientMessage::Complete {
+            query_instance_id: self.id,
+        };
+        let mut client_tx = self.client_tx.clone();
+        tokio::spawn(async move {
+            if let Err(e) = client_tx.send(message).await {
+                eprintln!("Failed to send complete message: {}", e);
+            }
+        });
+    }
+}
+
+pub struct Client<M: Memoize> {
+    tx: mpsc::Sender<ClientMessage>,
+    handle: tokio::task::JoinHandle<M>,
+}
+
+impl<M: Memoize> Client<M> {
+    pub fn new(tx: mpsc::Sender<ClientMessage>, handle: tokio::task::JoinHandle<M>) -> Self {
+        Self { tx, handle }
+    }
+
+    pub async fn create_query_instance(
+        &mut self,
+        logical_plan: LogicalPlan,
+    ) -> Result<QueryInstance, Error> {
+        let (id_tx, id_rx) = oneshot::channel();
+        // Need this to buffer > 0 so we do not block the optimizer.
+        let (physical_plan_tx, physical_plan_rx) = mpsc::channel(1);
+        let message = ClientMessage::Init {
+            logical_plan: logical_plan.clone(),
+            physical_plan_tx,
+            id_tx,
+        };
+        self.tx.send(message).await?;
+        let id = id_rx.await?;
+
+        Ok(QueryInstance {
+            id,
+            logical_plan,
+            physical_plan_rx,
+            client_tx: self.tx.clone(),
+        })
+    }
+
+    /// Shutdown the optimizer, returning the memo table.
+    pub async fn shutdown(mut self) -> Result<M, Error> {
+        self.tx.send(ClientMessage::Shutdown).await?;
+        let memo = self.handle.await.map_err(|_| Error::Placeholder)?;
+        Ok(memo)
+    }
+}
+
+pub enum ClientMessage {
+    /// Message to initiate optimization of a query instance.
+    Init {
+        /// The logical plan to be ingested.
+        logical_plan: LogicalPlan,
+
+        physical_plan_tx: mpsc::Sender<PhysicalPlan>,
+        /// The channel to send the query instance ID back to the caller.
+        id_tx: oneshot::Sender<QueryInstanceId>,
+    },
+    /// Message to indicate that the client is done with the query instance.
+    Complete {
+        /// The query instance ID to be completed.
+        query_instance_id: QueryInstanceId,
+    },
+    Shutdown,
+}

optd/src/core/optimizer/egest.rs

@@ -0,0 +1,203 @@
+use std::sync::Arc;
+
+use async_recursion::async_recursion;
+use futures::future::try_join_all;
+
+use crate::{
+    core::cir::{
+        Child, GoalMemberId, Operator, PartialPhysicalPlan, PhysicalExpressionId, PhysicalPlan,
+    },
+    core::error::Error,
+    memo::Memoize,
+};
+
+use super::Optimizer;
+
+impl<M: Memoize> Optimizer<M> {
+    /// Recursively transforms a physical expression ID in the memo into a complete physical plan.
+    ///
+    /// This function retrieves the physical expression from the memo and recursively
+    /// transforms any child goal members into their corresponding best physical plans.
+    ///
+    /// # Parameters
+    /// * `expression_id` - ID of the physical expression to transform into a complete plan.
+    ///
+    /// # Returns
+    /// * `Ok(Some(PhysicalPlan))` if all child plans were successfully constructed from their IDs.
+    /// * `Ok(None)` if any goal ID lacks a best expression ID.
+    /// * `Err(Error)` if a memo operation fails.
+    #[async_recursion]
+    pub(super) async fn egest_best_plan(
+        &self,
+        expression_id: PhysicalExpressionId,
+    ) -> Result<Option<PhysicalPlan>, Error> {
+        let expression = self.memo.materialize_physical_expr(expression_id).await?;
+
+        let child_results = try_join_all(
+            expression
+                .children
+                .iter()
+                .map(|child| self.egest_child_plan(child)),
+        )
+        .await?;
+
+        let child_plans = match child_results.into_iter().collect::<Option<Vec<_>>>() {
+            Some(plans) => plans,
+            None => return Ok(None),
+        };
+
+        Ok(Some(PhysicalPlan(Operator {
+            tag: expression.tag,
+            data: expression.data,
+            children: child_plans,
+        })))
+    }
+
+    /// Converts a physical expression ID to a partial physical plan.
+    ///
+    /// This method materializes the expression and recursively processes its children,
+    /// preserving goal references as unmaterialized plans.
+    ///
+    /// # Parameters
+    /// * `expression_id` - ID of the physical expression to convert to a partial plan.
+    ///
+    /// # Returns
+    /// * `PartialPhysicalPlan` - The materialized partial plan.
+    /// * `Err(Error)` if a memo operation fails.
+    pub(super) async fn egest_partial_plan(
+        &self,
+        expression_id: PhysicalExpressionId,
+    ) -> Result<PartialPhysicalPlan, Error> {
+        let expression = self.memo.materialize_physical_expr(expression_id).await?;
+
+        let children = try_join_all(
+            expression
+                .children
+                .iter()
+                .map(|child| self.egest_partial_child(child.clone())),
+        )
+        .await?;
+
+        Ok(PartialPhysicalPlan::Materialized(Operator {
+            tag: expression.tag,
+            data: expression.data,
+            children,
+        }))
+    }
+
+    async fn egest_child_plan(
+        &self,
+        child: &Child<GoalMemberId>,
+    ) -> Result<Option<Child<Arc<PhysicalPlan>>>, Error> {
+        match child {
+            Child::Singleton(member) => {
+                let plan = match self.process_goal_member(*member).await? {
+                    Some(plan) => plan,
+                    None => return Ok(None),
+                };
+                Ok(Some(Child::Singleton(plan.into())))
+            }
+            Child::VarLength(members) => {
+                let futures = members.iter().map(|member| async move {
+                    let plan = match self.process_goal_member(*member).await? {
+                        Some(plan) => plan,
+                        None => return Ok::<_, Error>(None),
+                    };
+                    Ok(Some(plan.into()))
+                });
+
+                let result_plans = match try_join_all(futures).await?.into_iter().collect() {
+                    Some(plans) => plans,
+                    None => return Ok(None),
+                };
+
+                Ok(Some(Child::VarLength(result_plans)))
+            }
+        }
+    }
+
+    async fn process_goal_member(
+        &self,
+        member: GoalMemberId,
+    ) -> Result<Option<PhysicalPlan>, Error> {
+        match member {
+            GoalMemberId::PhysicalExpressionId(expr_id) => self.egest_best_plan(expr_id).await,
+            GoalMemberId::GoalId(goal_id) => {
+                let (best_expr_id, _) =
+                    match self.memo.get_best_optimized_physical_expr(goal_id).await? {
+                        Some(expr) => expr,
+                        None => return Ok(None),
+                    };
+
+                self.egest_best_plan(best_expr_id).await
+            }
+        }
+    }
+
+    async fn egest_partial_child(
+        &self,
+        child: Child<GoalMemberId>,
+    ) -> Result<Child<Arc<PartialPhysicalPlan>>, Error> {
+        match child {
+            Child::Singleton(member) => match member {
+                GoalMemberId::GoalId(goal_id) => {
+                    let goal = self.memo.materialize_goal(goal_id).await?;
+                    Ok(Child::Singleton(
+                        PartialPhysicalPlan::UnMaterialized(goal).into(),
+                    ))
+                }
+                GoalMemberId::PhysicalExpressionId(expr_id) => {
+                    let expr = self.memo.materialize_physical_expr(expr_id).await?;
+
+                    let children = try_join_all(
+                        expr.children
+                            .iter()
+                            .map(|child| self.egest_partial_child(child.clone())),
+                    )
+                    .await?;
+
+                    let op = Operator {
+                        tag: expr.tag,
+                        data: expr.data,
+                        children,
+                    };
+
+                    Ok(Child::Singleton(
+                        PartialPhysicalPlan::Materialized(op).into(),
+                    ))
+                }
+            },
+            Child::VarLength(members) => {
+                let goals = try_join_all(members.into_iter().map(|member| async move {
+                    match member {
+                        GoalMemberId::GoalId(goal_id) => {
+                            let goal = self.memo.materialize_goal(goal_id).await?;
+                            Ok::<_, Error>(PartialPhysicalPlan::UnMaterialized(goal).into())
+                        }
+                        GoalMemberId::PhysicalExpressionId(expr_id) => {
+                            let expr = self.memo.materialize_physical_expr(expr_id).await?;
+
+                            let children = try_join_all(
+                                expr.children
+                                    .iter()
+                                    .map(|child| self.egest_partial_child(child.clone())),
+                            )
+                            .await?;
+
+                            let op = Operator {
+                                tag: expr.tag,
+                                data: expr.data,
+                                children,
+                            };
+
+                            Ok(PartialPhysicalPlan::Materialized(op).into())
+                        }
+                    }
+                }))
+                .await?;
+
+                Ok(Child::VarLength(goals))
+            }
+        }
+    }
+}