Try new structure

Author: alamb

src/llm/plan.rs renamed to src/llm/exec.rs

@@ -21,7 +21,6 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
-use crate::llm::expr::LLMExpr;
 
 /// This structure evaluates  a set of  expressions on a record
 /// batch producing a new record batch
@@ -33,17 +32,16 @@ pub struct LLMExec {
     /// Cache holding plan properties like equivalences, output partitioning etc.
     cache: PlanProperties,
     metrics: ExecutionPlanMetricsSet,
-    exprs: Vec<LLMExpr>,
+    // which expressions to evaluate
 }
 
 impl LLMExec {
-    pub fn new(input: Arc<dyn ExecutionPlan>, metrics: ExecutionPlanMetricsSet, exprs: Vec<LLMExpr>) -> Self {
+    pub fn new(input: Arc<dyn ExecutionPlan>, metrics: ExecutionPlanMetricsSet) -> Self {
         let cache = LLMExec::compute_properties(&input, input.schema().clone()).unwrap();
         Self {
             input,
             cache,
             metrics,
-            exprs
         }
     }
 
@@ -171,7 +169,6 @@ impl LLMStream {
     }
 }
 
-
 // TODO:
 //  pretty batch and concat with promotion
 //  as the input of the LLM function

src/llm/expr.rs

@@ -1,5 +1,3 @@
+pub mod exec;
 pub mod functions;
-pub mod logical;
-pub mod physical_planner;
-pub mod plan;
-mod expr;
+pub mod physical_optimizer;

src/llm/logical.rs

@@ -0,0 +1,46 @@
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::ExecutionPlan;
+use std::sync::Arc;
+
+#[derive(Debug)]
+pub struct LLMRule {}
+
+impl PhysicalOptimizerRule for LLMRule {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        config: &ConfigOptions,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        /// Insert a LLMExec node in front of this projection if there are any llm functions in it
+        ///
+        /// For example, if the projection is:
+        /// ```text
+        ///   ProjectionExec(["A", "B", llm_func('foo', "C") + 1])
+        /// ```
+        /// Rewrite to
+        ///   ProjectionExec(["A", "B", "llm_func('foo', "C")" + 1]) <-- note here that the llm function is not evaluated and instead is a new column
+        ///     LLMExec(["A", "B", llm_func('foo', "C")])
+        ///
+        /// This is similar to the rewrite done by the CommonSubexprEliminationRule
+
+        /// replace ProjectionExec with LLMExec if there are any llm functions
+        let Some(proj_exec) = plan.as_any().downcast_ref::<ProjectionExec>() else {
+            return Ok(plan);
+        };
+
+        /// find llm functions in the expressions
+        //proj_exec.expr()
+        todo!()
+    }
+
+    fn name(&self) -> &str {
+        "llm"
+    }
+
+    /// verify the schema has not changed
+    fn schema_check(&self) -> bool {
+        true
+    }
+}