[FEAT] (ACTORS-2) Add optimization pass to split Project into ActorPo…

…olProject (#2627) Adds logical optimization for splitting a `Project` with stateful UDFs into a chain of `(Project -> ActorPool -> ActorPool -> ... -> Project)` We then rely on ProjectionPushdown rules to optimize this chain (applying column pruning, projection fusion etc). --------- Co-authored-by: Jay Chia <[email protected]@users.noreply.github.com> Co-authored-by: Desmond Cheong <[email protected]>
Eventual-Inc · Aug 23, 2024 · 9df6beb · 9df6beb
1 parent a235f80
commit 9df6beb
Show file tree

Hide file tree

Showing 10 changed files with 1,400 additions and 36 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -161,6 +161,7 @@ sketches-ddsketch = {version = "0.2.2", features = ["use_serde"]}
 snafu = {version = "0.7.4", features = ["futures"]}
 sqlparser = "0.49.0"
 sysinfo = "0.30.12"
+test-log = "0.2.16"
 tiktoken-rs = "0.5.9"
 tokio = {version = "1.37.0", features = [
   "net",

diff --git a/daft/udf.py b/daft/udf.py
@@ -445,6 +445,7 @@ def udf(
     num_gpus: float | None = None,
     memory_bytes: int | None = None,
     batch_size: int | None = None,
+    _concurrency: int | None = None,
 ) -> Callable[[UserProvidedPythonFunction | type], StatelessUDF | StatefulUDF]:
     """Decorator to convert a Python function into a UDF
 

diff --git a/src/daft-dsl/src/expr.rs b/src/daft-dsl/src/expr.rs
@@ -610,7 +610,7 @@ impl Expr {
         }
     }
 
-    pub(super) fn with_new_children(&self, children: Vec<ExprRef>) -> Expr {
+    pub fn with_new_children(&self, children: Vec<ExprRef>) -> Expr {
         use Expr::*;
         match self {
             // no children

diff --git a/src/daft-plan/Cargo.toml b/src/daft-plan/Cargo.toml
@@ -36,6 +36,7 @@ snafu = {workspace = true}
 daft-functions = {path = "../daft-functions", default-features = false}
 pretty_assertions = {workspace = true}
 rstest = {workspace = true}
+test-log = {workspace = true}
 
 [features]
 python = [